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 13, 10:45 UTC
[Core] Encoder separation for Encode-Prefill-Decode Disaggregation (#25233)
Signed-off-by: n00909098 <nguyen.kha.long@huawei.com> Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com> Signed-off-by: herotai214 <herotai214@gmail.com> Signed-off-by: Khuong Le <khuong.le.manh@huawei.com> Signed-off-by: Khuong Le <lemanhkhuong2611@gmail.com> Co-authored-by: n00909098 <nguyen.kha.long@huawei.com> Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com> Co-authored-by: herotai214 <herotai214@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Khuong Le <khuong.le.manh@huawei.com> Co-authored-by: Khuong Le <lemanhkhuong2611@gmail.com>
Chenguang Zheng committed
Nov 12, 2025 at 10:58 UTC
4ccffe561f0bc9fe007fdaafe8342171592db7db
31 files changed
+5025
-41
docs/assets/features/disagg_encoder/disagg_encoder_flow.png
Binary files /dev/null and b/docs/assets/features/disagg_encoder/disagg_encoder_flow.png differ
docs/features/disagg_encoder.md
new
+75
@@ -0,0 +1,75 @@
1
+# Disaggregated Encoder
2
+
3
+A **disaggregated encoder** runs the vision-encoder stage of a multimodal LLM in a process that is separate from the pre-fill / decoder stage. Deploying these two stages in independent vLLM instances brings three practical benefits:
4
+
5
+1. **Independent, fine-grained scaling**
6
+2. **Lower time-to-first-token (TTFT)**
7
+3. **Cross-process reuse and caching of encoder outputs**
8
+
9
+Design doc: <https://docs.google.com/document/d/1aed8KtC6XkXtdoV87pWT0a8OJlZ-CpnuLLzmR8l9BAE>
10
+
11
+---
12
+
13
+## 1 Motivation
14
+
15
+### 1. Independent, fine-grained scaling
16
+
17
+* Vision encoders are lightweight, while language models are orders of magnitude larger.
18
+* The language model can be parallelised without affecting the encoder fleet.
19
+* Encoder nodes can be added or removed independently.
20
+
21
+### 2. Lower time-to-first-token (TTFT)
22
+
23
+* Language-only requests bypass the vision encoder entirely.
24
+* Encoder output is injected only at required attention layers, shortening the pre-fill critical path.
25
+
26
+### 3. Cross-process reuse and caching
27
+
28
+* In-process encoders confine reuse to a single worker.
29
+* A remote, shared cache lets any worker retrieve existing embeddings, eliminating redundant computation.
30
+
31
+---
32
+
33
+## 2 Usage Example
34
+
35
+The current reference pathway is **SharedStorageConnector**.
36
+Below ready-to-run scripts shows the workflow:
37
+
38
+1 Encoder instance + 1 PD instance:
39
+`examples/online_serving/disaggregated_encoder/shared_storage_connector/disagg_encoder_example.sh`
40
+
41
+1 Encoder instance + 1 Prefill instance + 1 Decode instance:
42
+`examples/online_serving/disaggregated_encoder/shared_storage_connector/disagg_epd_example.sh`
43
+
44
+---
45
+
46
+## 3 Test Script
47
+
48
+Please refer to the directories `tests/v1/ec_connector`
49
+
50
+## 4 Development
51
+
52
+Disaggregated encoding is implemented by running two parts:
53
+
54
+* **Encoder instance** – a vLLM instance to performs vision encoding.
55
+* **Prefill/Decode (PD) instance(s)** – runs language pre-fill and decode.
56
+ * PD can be in either a single normal instance with `disagg_encoder_example.sh` (E->PD) or in disaggregated instances with `disagg_epd_example.sh` (E->P->D)
57
+
58
+A connector transfers encoder-cache (EC) embeddings from the encoder instance to the PD instance.
59
+All related code is under `vllm/distributed/ec_transfer`.
60
+
61
+### Key abstractions
62
+
63
+* **ECConnector** – interface for retrieving EC caches produced by the encoder.
64
+ * *Scheduler role* – checks cache existence and schedules loads.
65
+ * *Worker role* – loads the embeddings into memory.
66
+
67
+Here is a figure illustrating disaggregate encoder flow:
68
+
69
+
70
+
71
+For the PD disaggregation part, the Prefill instance receive cache exactly the same as the disaggregate encoder flow above. Prefill instance executes 1 step (prefill -> 1 token output) and then transfer KV cache to the Decode instance for the remaining execution. The KV transfer part purely happens after the execute of the PDinstance.
72
+
73
+`docs/features/disagg_prefill.md` shows the brief idea about the disaggregated prefill (v0)
74
+
75
+We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D;
examples/online_serving/disaggregated_encoder/README.md
new
+119
@@ -0,0 +1,119 @@
1
+# Disaggregated Encoder
2
+
3
+These example scripts that demonstrate the disaggregated encoder (EPD) features of vLLM.
4
+
5
+For a detailed explanation of the EPD features, please refer to the [Disaggregated Encoder Feature Documentation](../../../docs/features/disagg_encoder.md).
6
+
7
+## Files
8
+
9
+- `disagg_epd_proxy.py` - Proxy script that demonstrates the XeYpZd setup (X encode instances, Y prefill instances, Z decode instances). Currently stable for the 1e1p1d configuration.
10
+
11
+- `disagg_1e1p1d_example.sh` - Sets up the 1e1p1d configuration, runs the VisionArena benchmark, and processes a single request with a local image.
12
+
13
+- `disagg_1e1pd_example.sh` - Sets up the 1e1pd configuration, runs the VisionArena benchmark, and processes a single request with a local image.
14
+
15
+### Custom Configuration
16
+
17
+```bash
18
+# Use specific GPUs
19
+GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash disagg_1e1p1d_example.sh
20
+
21
+# Use specific ports
22
+ENDPOINT_PORT=10001 bash disagg_1e1p1d_example.sh
23
+
24
+# Use specific model
25
+MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash disagg_1e1p1d_example.sh
26
+
27
+# Use specific storage path
28
+EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash disagg_1e1p1d_example.sh
29
+```
30
+
31
+## Encoder Instances
32
+
33
+Encoder engines should be launched with the following flags:
34
+
35
+- `--enforce-eager` **(required)** – The current EPD implementation is only compatible with encoder instances running in this mode.
36
+
37
+- `--no-enable-prefix-caching` **(required)** – Encoder instances do not consume KV cache; prefix caching is disabled to avoid conflicts with other features.
38
+
39
+- `--max-num-batched-tokens=<large value>` **(default: 2048)** – This flag controls the token scheduling budget per decoding step and is irrelevant to encoder-only instances. **Set it to a very high value (effectively unlimited) to bypass scheduler limitations.** The actual token budget is managed by the encoder cache manager.
40
+
41
+## Local media inputs
42
+
43
+To support local image inputs (from your ```MEDIA_PATH``` directory), add the following flag to the encoder instance:
44
+
45
+```bash
46
+--allowed-local-media-path $MEDIA_PATH
47
+```
48
+
49
+The vllm instances and `disagg_encoder_proxy` supports local URIs with ```{"url": "file://'"$MEDIA_PATH_FILENAME"'}``` as multimodal inputs. Each URI is passed unchanged from the `disagg_encoder_proxy` to the encoder instance so that the encoder can load the media locally.
50
+
51
+## EC connector and KV transfer
52
+
53
+The `ECSharedStorageConnector` is used to store the encoder cache on local disk and facilitate transfer. To enable the encoder disaggregation feature, add the following configuration:
54
+
55
+```bash
56
+# Add to encoder instance:
57
+--ec-transfer-config '{
58
+ "ec_connector": "ECSharedStorageConnector",
59
+ "ec_role": "ec_producer",
60
+ "ec_connector_extra_config": {
61
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
62
+ }
63
+}'
64
+
65
+# Add to prefill/prefill+decode instance:
66
+--ec-transfer-config '{
67
+ "ec_connector": "ECSharedStorageConnector",
68
+ "ec_role": "ec_consumer",
69
+ "ec_connector_extra_config": {
70
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
71
+ }
72
+}'
73
+```
74
+
75
+`$EC_SHARED_STORAGE_PATH` is the path where the EC connector temporarily stores the cache.
76
+
77
+If you enable prefill instance (`--prefill-servers-urls` not disabled), you will need --kv-transfer-config to facilitate the PD disaggregation. Currently, we use the `NixlConnector` for this purpose. Refer to `tests/v1/kv_connector/nixl_integration` for more example codes on PD disaggregation with Nixl.
78
+
79
+```bash
80
+# Add to prefill instance:
81
+--kv-transfer-config '{
82
+ "kv_connector": "NixlConnector",
83
+ "kv_role": "kv_producer"
84
+}'
85
+
86
+# Add to decode instance:
87
+--kv-transfer-config '{
88
+ "kv_connector": "NixlConnector",
89
+ "kv_role": "kv_consumer"
90
+}'
91
+```
92
+
93
+## Proxy Instance Flags (`disagg_epd_proxy.py`)
94
+
95
+| Flag | Description |
96
+|------|-------------|
97
+| `--encode-servers-urls` | Comma-separated list of encoder endpoints. Every multimodal item extracted from the request is fanned out to one of these URLs in a round-robin fashion. |
98
+| `--prefill-servers-urls` | Comma-separated list of prefill endpoints. Set to `disable`, `none`, or `""` to skip the dedicated prefill phase and run E+PD (encoder + combined prefill/decode). |
99
+| `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. |
100
+| `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). |
101
+
102
+Example usage:
103
+For E + PD setup:
104
+
105
+```bash
106
+$ python disagg_encoder_proxy.py \
107
+ --encode-servers-urls "http://e1:8001,http://e2:8002" \
108
+ --prefill-servers-urls "disable" \
109
+ --decode-servers-urls "http://pd1:8003,http://pd2:8004"
110
+```
111
+
112
+For E + P + D setup:
113
+
114
+```bash
115
+$ python disagg_encoder_proxy.py \
116
+ --encode-servers-urls "http://e1:8001,http://e2:8001" \
117
+ --prefill-servers-urls "http://p1:8003,http://p2:8004" \
118
+ --decode-servers-urls "http://d1:8005,http://d2:8006"
119
+```
examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh
new
+221
@@ -0,0 +1,221 @@
1
+#!/bin/bash
2
+set -euo pipefail
3
+
4
+declare -a PIDS=()
5
+
6
+###############################################################################
7
+# Configuration -- override via env before running
8
+###############################################################################
9
+MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
10
+LOG_PATH="${LOG_PATH:-./logs}"
11
+mkdir -p $LOG_PATH
12
+
13
+ENCODE_PORT="${ENCODE_PORT:-19534}"
14
+PREFILL_PORT="${PREFILL_PORT:-19535}"
15
+DECODE_PORT="${DECODE_PORT:-19536}"
16
+PROXY_PORT="${PROXY_PORT:-10001}"
17
+
18
+GPU_E="${GPU_E:-2}"
19
+GPU_P="${GPU_P:-2}"
20
+GPU_D="${GPU_D:-3}"
21
+
22
+EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
23
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}" # wait_for_server timeout
24
+
25
+NUM_PROMPTS="${NUM_PROMPTS:-100}" # number of prompts to send in benchmark
26
+
27
+export UCX_TLS=all
28
+export UCX_NET_DEVICES=all
29
+
30
+###############################################################################
31
+# Helpers
32
+###############################################################################
33
+# Find the git repository root directory
34
+GIT_ROOT=$(git rev-parse --show-toplevel)
35
+
36
+START_TIME=$(date +"%Y%m%d_%H%M%S")
37
+ENC_LOG=$LOG_PATH/encoder_${START_TIME}.log
38
+P_LOG=$LOG_PATH/p_${START_TIME}.log
39
+D_LOG=$LOG_PATH/d_${START_TIME}.log
40
+PROXY_LOG=$LOG_PATH/proxy_${START_TIME}.log
41
+
42
+wait_for_server() {
43
+ local port=$1
44
+ timeout "$TIMEOUT_SECONDS" bash -c "
45
+ until curl -s localhost:$port/v1/chat/completions > /dev/null; do
46
+ sleep 1
47
+ done" && return 0 || return 1
48
+}
49
+
50
+# Cleanup function
51
+cleanup() {
52
+ echo "Stopping everything…"
53
+ trap - INT TERM USR1 # prevent re-entrancy
54
+
55
+ # Kill all tracked PIDs
56
+ for pid in "${PIDS[@]}"; do
57
+ if kill -0 "$pid" 2>/dev/null; then
58
+ echo "Killing process $pid"
59
+ kill "$pid" 2>/dev/null
60
+ fi
61
+ done
62
+
63
+ # Wait a moment for graceful shutdown
64
+ sleep 2
65
+
66
+ # Force kill any remaining processes
67
+ for pid in "${PIDS[@]}"; do
68
+ if kill -0 "$pid" 2>/dev/null; then
69
+ echo "Force killing process $pid"
70
+ kill -9 "$pid" 2>/dev/null
71
+ fi
72
+ done
73
+
74
+ # Kill the entire process group as backup
75
+ kill -- -$$ 2>/dev/null
76
+
77
+ echo "All processes stopped."
78
+ exit 0
79
+}
80
+
81
+trap cleanup INT
82
+trap cleanup USR1
83
+trap cleanup TERM
84
+
85
+# clear previous cache
86
+echo "remove previous ec cache folder"
87
+rm -rf $EC_SHARED_STORAGE_PATH
88
+
89
+echo "make ec cache folder"
90
+mkdir -p $EC_SHARED_STORAGE_PATH
91
+
92
+###############################################################################
93
+# Encoder worker
94
+###############################################################################
95
+CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
96
+ --gpu-memory-utilization 0.01 \
97
+ --port "$ENCODE_PORT" \
98
+ --enforce-eager \
99
+ --enable-request-id-headers \
100
+ --no-enable-prefix-caching \
101
+ --max-num-batched-tokens 114688 \
102
+ --max-num-seqs 128 \
103
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
104
+ --ec-transfer-config '{
105
+ "ec_connector": "ECSharedStorageConnector",
106
+ "ec_role": "ec_producer",
107
+ "ec_connector_extra_config": {
108
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
109
+ }
110
+ }' \
111
+ >"${ENC_LOG}" 2>&1 &
112
+
113
+PIDS+=($!)
114
+
115
+###############################################################################
116
+# Prefill worker
117
+###############################################################################
118
+CUDA_VISIBLE_DEVICES="$GPU_P" \
119
+UCX_NET_DEVICES=all \
120
+VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
121
+vllm serve "$MODEL" \
122
+ --gpu-memory-utilization 0.7 \
123
+ --port "$PREFILL_PORT" \
124
+ --enforce-eager \
125
+ --enable-request-id-headers \
126
+ --max-num-seqs 128 \
127
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
128
+ --ec-transfer-config '{
129
+ "ec_connector": "ECSharedStorageConnector",
130
+ "ec_role": "ec_consumer",
131
+ "ec_connector_extra_config": {
132
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
133
+ }
134
+ }' \
135
+ --kv-transfer-config '{
136
+ "kv_connector": "NixlConnector",
137
+ "kv_role": "kv_producer"
138
+ }' \
139
+ >"${P_LOG}" 2>&1 &
140
+
141
+PIDS+=($!)
142
+
143
+###############################################################################
144
+# Decode worker
145
+###############################################################################
146
+CUDA_VISIBLE_DEVICES="$GPU_D" \
147
+UCX_NET_DEVICES=all \
148
+VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
149
+vllm serve "$MODEL" \
150
+ --gpu-memory-utilization 0.7 \
151
+ --port "$DECODE_PORT" \
152
+ --enforce-eager \
153
+ --enable-request-id-headers \
154
+ --max-num-seqs 128 \
155
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
156
+ --kv-transfer-config '{
157
+ "kv_connector": "NixlConnector",
158
+ "kv_role": "kv_consumer"
159
+ }' \
160
+ >"${D_LOG}" 2>&1 &
161
+
162
+PIDS+=($!)
163
+
164
+# Wait for workers
165
+wait_for_server $ENCODE_PORT
166
+wait_for_server $PREFILL_PORT
167
+wait_for_server $DECODE_PORT
168
+
169
+###############################################################################
170
+# Proxy
171
+###############################################################################
172
+python disagg_epd_proxy.py \
173
+ --host "0.0.0.0" \
174
+ --port "$PROXY_PORT" \
175
+ --encode-servers-urls "http://localhost:$ENCODE_PORT" \
176
+ --prefill-servers-urls "http://localhost:$PREFILL_PORT" \
177
+ --decode-servers-urls "http://localhost:$DECODE_PORT" \
178
+ >"${PROXY_LOG}" 2>&1 &
179
+
180
+PIDS+=($!)
181
+
182
+wait_for_server $PROXY_PORT
183
+echo "All services are up!"
184
+
185
+###############################################################################
186
+# Benchmark
187
+###############################################################################
188
+echo "Running benchmark (stream)..."
189
+vllm bench serve \
190
+ --model $MODEL \
191
+ --backend openai-chat \
192
+ --endpoint /v1/chat/completions \
193
+ --dataset-name hf \
194
+ --dataset-path lmarena-ai/VisionArena-Chat \
195
+ --seed 0 \
196
+ --num-prompts $NUM_PROMPTS \
197
+ --port $PROXY_PORT
198
+
199
+PIDS+=($!)
200
+
201
+###############################################################################
202
+# Single request with local image
203
+###############################################################################
204
+echo "Running single request with local image (non-stream)..."
205
+curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
206
+ -H "Content-Type: application/json" \
207
+ -d '{
208
+ "model": "'${MODEL}'",
209
+ "messages": [
210
+ {"role": "system", "content": "You are a helpful assistant."},
211
+ {"role": "user", "content": [
212
+ {"type": "image_url", "image_url": {"url": "file://'"${GIT_ROOT}"'/tests/v1/ec_connector/integration/hato.jpg"}},
213
+ {"type": "text", "text": "What is in this image?"}
214
+ ]}
215
+ ]
216
+ }'
217
+
218
+
219
+# cleanup
220
+echo "cleanup..."
221
+cleanup
\ No newline at end of file
examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh
new
+186
@@ -0,0 +1,186 @@
1
+#!/bin/bash
2
+set -euo pipefail
3
+
4
+declare -a PIDS=()
5
+
6
+###############################################################################
7
+# Configuration -- override via env before running
8
+###############################################################################
9
+MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
10
+LOG_PATH="${LOG_PATH:-./logs}"
11
+mkdir -p $LOG_PATH
12
+
13
+ENCODE_PORT="${ENCODE_PORT:-19534}"
14
+PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19535}"
15
+PROXY_PORT="${PROXY_PORT:-10001}"
16
+
17
+GPU_E="${GPU_E:-0}"
18
+GPU_PD="${GPU_PD:-1}"
19
+
20
+EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
21
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}" # wait_for_server timeout
22
+
23
+NUM_PROMPTS="${NUM_PROMPTS:-100}" # number of prompts to send in benchmark
24
+
25
+###############################################################################
26
+# Helpers
27
+###############################################################################
28
+# Find the git repository root directory
29
+GIT_ROOT=$(git rev-parse --show-toplevel)
30
+
31
+START_TIME=$(date +"%Y%m%d_%H%M%S")
32
+ENC_LOG=$LOG_PATH/encoder_${START_TIME}.log
33
+PD_LOG=$LOG_PATH/pd_${START_TIME}.log
34
+PROXY_LOG=$LOG_PATH/proxy_${START_TIME}.log
35
+
36
+wait_for_server() {
37
+ local port=$1
38
+ timeout "$TIMEOUT_SECONDS" bash -c "
39
+ until curl -s localhost:$port/v1/chat/completions > /dev/null; do
40
+ sleep 1
41
+ done" && return 0 || return 1
42
+}
43
+
44
+# Cleanup function
45
+cleanup() {
46
+ echo "Stopping everything…"
47
+ trap - INT TERM USR1 # prevent re-entrancy
48
+
49
+ # Kill all tracked PIDs
50
+ for pid in "${PIDS[@]}"; do
51
+ if kill -0 "$pid" 2>/dev/null; then
52
+ echo "Killing process $pid"
53
+ kill "$pid" 2>/dev/null
54
+ fi
55
+ done
56
+
57
+ # Wait a moment for graceful shutdown
58
+ sleep 2
59
+
60
+ # Force kill any remaining processes
61
+ for pid in "${PIDS[@]}"; do
62
+ if kill -0 "$pid" 2>/dev/null; then
63
+ echo "Force killing process $pid"
64
+ kill -9 "$pid" 2>/dev/null
65
+ fi
66
+ done
67
+
68
+ # Kill the entire process group as backup
69
+ kill -- -$$ 2>/dev/null
70
+
71
+ echo "All processes stopped."
72
+ exit 0
73
+}
74
+
75
+trap cleanup INT
76
+trap cleanup USR1
77
+trap cleanup TERM
78
+
79
+# clear previous cache
80
+echo "remove previous ec cache folder"
81
+rm -rf $EC_SHARED_STORAGE_PATH
82
+
83
+echo "make ec cache folder"
84
+mkdir -p $EC_SHARED_STORAGE_PATH
85
+
86
+###############################################################################
87
+# Encoder worker
88
+###############################################################################
89
+CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
90
+ --gpu-memory-utilization 0.01 \
91
+ --port "$ENCODE_PORT" \
92
+ --enforce-eager \
93
+ --enable-request-id-headers \
94
+ --no-enable-prefix-caching \
95
+ --max-num-batched-tokens 114688 \
96
+ --max-num-seqs 128 \
97
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
98
+ --ec-transfer-config '{
99
+ "ec_connector": "ECSharedStorageConnector",
100
+ "ec_role": "ec_producer",
101
+ "ec_connector_extra_config": {
102
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
103
+ }
104
+ }' \
105
+ >"${ENC_LOG}" 2>&1 &
106
+
107
+PIDS+=($!)
108
+
109
+###############################################################################
110
+# Prefill+Decode worker
111
+###############################################################################
112
+CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
113
+ --gpu-memory-utilization 0.7 \
114
+ --port "$PREFILL_DECODE_PORT" \
115
+ --enforce-eager \
116
+ --enable-request-id-headers \
117
+ --max-num-seqs 128 \
118
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
119
+ --ec-transfer-config '{
120
+ "ec_connector": "ECSharedStorageConnector",
121
+ "ec_role": "ec_consumer",
122
+ "ec_connector_extra_config": {
123
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
124
+ }
125
+ }' \
126
+ >"${PD_LOG}" 2>&1 &
127
+
128
+PIDS+=($!)
129
+
130
+# Wait for workers
131
+wait_for_server $ENCODE_PORT
132
+wait_for_server $PREFILL_DECODE_PORT
133
+
134
+###############################################################################
135
+# Proxy
136
+###############################################################################
137
+python disagg_epd_proxy.py \
138
+ --host "0.0.0.0" \
139
+ --port "$PROXY_PORT" \
140
+ --encode-servers-urls "http://localhost:$ENCODE_PORT" \
141
+ --prefill-servers-urls "disable" \
142
+ --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
143
+ >"${PROXY_LOG}" 2>&1 &
144
+
145
+PIDS+=($!)
146
+
147
+wait_for_server $PROXY_PORT
148
+echo "All services are up!"
149
+
150
+###############################################################################
151
+# Benchmark
152
+###############################################################################
153
+echo "Running benchmark (stream)..."
154
+vllm bench serve \
155
+ --model $MODEL \
156
+ --backend openai-chat \
157
+ --endpoint /v1/chat/completions \
158
+ --dataset-name hf \
159
+ --dataset-path lmarena-ai/VisionArena-Chat \
160
+ --seed 0 \
161
+ --num-prompts $NUM_PROMPTS \
162
+ --port $PROXY_PORT
163
+
164
+PIDS+=($!)
165
+
166
+###############################################################################
167
+# Single request with local image
168
+###############################################################################
169
+echo "Running single request with local image (non-stream)..."
170
+curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
171
+ -H "Content-Type: application/json" \
172
+ -d '{
173
+ "model": "'${MODEL}'",
174
+ "messages": [
175
+ {"role": "system", "content": "You are a helpful assistant."},
176
+ {"role": "user", "content": [
177
+ {"type": "image_url", "image_url": {"url": "file://'"${GIT_ROOT}"'/tests/v1/ec_connector/integration/hato.jpg"}},
178
+ {"type": "text", "text": "What is in this image?"}
179
+ ]}
180
+ ]
181
+ }'
182
+
183
+
184
+# cleanup
185
+echo "cleanup..."
186
+cleanup
\ No newline at end of file
examples/online_serving/disaggregated_encoder/disagg_epd_proxy.py
new
+606
@@ -0,0 +1,606 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""
5
+disagg_encoder_proxy.py
6
+
7
+Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
8
+clusters:
9
+ • encode (multimodal feature extraction)
10
+ • decode (language-model inference)
11
+
12
+For MM input we:
13
+ 1. Extract *every* image/audio item.
14
+ 2. Fire N concurrent requests to the encoder cluster
15
+ (one request per item, with **all text removed**).
16
+ 3. Wait for all of them to succeed.
17
+ 4. Forward the *original* request to a decode server.
18
+"""
19
+
20
+from __future__ import annotations
21
+
22
+import argparse
23
+import asyncio
24
+import logging
25
+import os
26
+import random
27
+import uuid
28
+from collections.abc import AsyncIterator
29
+
30
+import aiohttp
31
+import uvicorn
32
+from fastapi import FastAPI, HTTPException, Request
33
+from fastapi.responses import JSONResponse, StreamingResponse
34
+
35
+###############################################################################
36
+# FastAPI app & global state
37
+###############################################################################
38
+
39
+logging.basicConfig(
40
+ level=logging.DEBUG, format="%(asctime)s %(levelname)s: %(message)s"
41
+)
42
+logger = logging.getLogger("proxy")
43
+
44
+app = FastAPI()
45
+encode_session: aiohttp.ClientSession | None = None
46
+prefill_session: aiohttp.ClientSession | None = None
47
+decode_session: aiohttp.ClientSession | None = None
48
+
49
+###############################################################################
50
+# Utils
51
+###############################################################################
52
+
53
+
54
+MM_TYPES = {"image_url", "audio_url", "input_audio"}
55
+
56
+
57
+def extract_mm_items(request_data: dict) -> list[dict]:
58
+ """
59
+ Return *all* image/audio items that appear anywhere in `messages`.
60
+
61
+ Each returned dict looks like:
62
+ { "type": "image_url", "image_url": {...} }
63
+ """
64
+ items: list[dict] = []
65
+ for msg in request_data.get("messages", []):
66
+ content = msg.get("content")
67
+ if not isinstance(content, list):
68
+ continue
69
+
70
+ for item in content:
71
+ if item.get("type") in MM_TYPES:
72
+ items.append(item)
73
+ return items
74
+
75
+
76
+async def fanout_encoder_primer(
77
+ orig_request: dict,
78
+ e_urls: list[str],
79
+ req_id: str,
80
+) -> None:
81
+ """
82
+ 1. Build one request *per MM item* with all text removed.
83
+ 2. Send them concurrently to the encode cluster.
84
+ 3. Raise if any of them fails.
85
+ """
86
+ logger.info("[%s] Processing multimodal items...", req_id)
87
+
88
+ mm_items = extract_mm_items(orig_request)
89
+ if not mm_items:
90
+ logger.info("[%s] No multimodal items, skipping encoder", req_id)
91
+ return # nothing to do
92
+
93
+ logger.info("[%s] got %d multimodal items...", req_id, len(mm_items))
94
+
95
+ tasks = []
96
+
97
+ # Round-robin over encode servers to distribute load a bit
98
+ url_cycle = (e_urls[i % len(e_urls)] for i in range(len(mm_items)))
99
+
100
+ for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
101
+ # Derive a *child* request id: <parent>:<index>:<random-short>
102
+ child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}"
103
+ headers = {"x-request-id": child_req_id}
104
+
105
+ encoder_req = {
106
+ # You *may* need to keep additional fields
107
+ "model": orig_request.get("model"),
108
+ "messages": [
109
+ {"role": "user", "content": [item]},
110
+ ],
111
+ # Only need 1 token so the server actually runs the encoder path
112
+ "max_tokens": 1,
113
+ "stream": False,
114
+ }
115
+ tasks.append(
116
+ encode_session.post(
117
+ f"{target_url}/v1/chat/completions",
118
+ json=encoder_req,
119
+ headers=headers,
120
+ )
121
+ )
122
+
123
+ results = await asyncio.gather(*tasks, return_exceptions=True)
124
+
125
+ # Fail fast if any sub-request failed
126
+ for idx, r in enumerate(results):
127
+ if isinstance(r, Exception):
128
+ logger.error(
129
+ "[%s] Encoder request #%d raised exception: %s",
130
+ req_id,
131
+ idx,
132
+ r,
133
+ exc_info=r,
134
+ )
135
+ raise HTTPException(
136
+ status_code=502, detail=f"Encoder request failed: {str(r)}"
137
+ )
138
+ if r.status != 200:
139
+ try:
140
+ detail = await r.text()
141
+ except Exception:
142
+ detail = "<unable to read body>"
143
+ logger.error(
144
+ "[%s] Encoder request #%d returned status %s: %s",
145
+ req_id,
146
+ idx,
147
+ r.status,
148
+ detail,
149
+ )
150
+ raise HTTPException(
151
+ status_code=r.status,
152
+ detail=f"Encoder request failed: {detail}",
153
+ )
154
+
155
+ logger.info(
156
+ "[%s] All %d encoder requests completed successfully", req_id, len(mm_items)
157
+ )
158
+
159
+
160
+async def maybe_prefill(
161
+ req_data: dict,
162
+ p_url: str,
163
+ req_id: str,
164
+) -> dict:
165
+ """
166
+ - Do prefill-only task if p_url exist;
167
+ - Return modified request data with kv transfer params (for nixl connector)
168
+ - Else, skip and return the original request data for decode
169
+ """
170
+ if p_url:
171
+ logger.info("[%s] Processing through prefill: %s", req_id, p_url)
172
+
173
+ prefill_response = await process_prefill_stage(req_data, p_url, req_id)
174
+ # for nixl connector to facilitate kv transfer...
175
+ prefill_response_json = await prefill_response.json()
176
+ kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
177
+ if kv_transfer_params:
178
+ req_data["kv_transfer_params"] = kv_transfer_params
179
+
180
+ return req_data
181
+ else:
182
+ return req_data
183
+
184
+
185
+async def process_prefill_stage(
186
+ req_data: dict,
187
+ p_url: str,
188
+ req_id: str,
189
+) -> dict:
190
+ """Process request through Prefill stage and return kv_transfer_params"""
191
+ logger.info("[%s] Sending prefill request to: %s", req_id, p_url)
192
+
193
+ prefill_request = req_data.copy()
194
+ prefill_request["kv_transfer_params"] = {
195
+ "do_remote_decode": True,
196
+ "do_remote_prefill": False,
197
+ "remote_engine_id": None,
198
+ "remote_block_ids": None,
199
+ "remote_host": None,
200
+ "remote_port": None,
201
+ }
202
+ prefill_request["stream"] = False
203
+ prefill_request["max_tokens"] = 1
204
+ if "max_completion_tokens" in prefill_request:
205
+ prefill_request["max_completion_tokens"] = 1
206
+ if "stream_options" in prefill_request:
207
+ del prefill_request["stream_options"]
208
+
209
+ headers = {"x-request-id": req_id}
210
+ try:
211
+ prefill_response = await prefill_session.post(
212
+ f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers
213
+ )
214
+ prefill_response.raise_for_status()
215
+
216
+ if prefill_response.status != 200:
217
+ error_text = await prefill_response.text()
218
+ logger.error(
219
+ "[%s] Prefill request failed with status %d: %s",
220
+ req_id,
221
+ prefill_response.status,
222
+ error_text,
223
+ )
224
+ raise HTTPException(
225
+ status_code=prefill_response.status,
226
+ detail={"error": "Prefill request failed", "message": error_text},
227
+ )
228
+ logger.info("[%s] Prefill request completed successfully", req_id)
229
+
230
+ return prefill_response
231
+
232
+ except Exception as e:
233
+ logger.error("Prefill processing failed: %s", str(e))
234
+ raise HTTPException(
235
+ status_code=500,
236
+ detail={"error": "Prefill processing error", "message": str(e)},
237
+ ) from e
238
+
239
+
240
+###############################################################################
241
+# Middleware for request/response logging
242
+###############################################################################
243
+
244
+
245
+@app.middleware("http")
246
+async def log_requests(request: Request, call_next):
247
+ """Middleware to log all incoming requests and responses"""
248
+ req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
249
+
250
+ # Log incoming request
251
+ logger.info(
252
+ ">>> [%s] %s %s from %s",
253
+ req_id,
254
+ request.method,
255
+ request.url.path,
256
+ request.client.host if request.client else "unknown",
257
+ )
258
+
259
+ try:
260
+ # Process request
261
+ response = await call_next(request)
262
+
263
+ # Log response
264
+ logger.info(
265
+ "<<< [%s] %s %s completed with status %d",
266
+ req_id,
267
+ request.method,
268
+ request.url.path,
269
+ response.status_code,
270
+ )
271
+
272
+ return response
273
+ except Exception as e:
274
+ # Log errors
275
+ logger.exception(
276
+ "!!! [%s] %s %s failed with error: %s",
277
+ req_id,
278
+ request.method,
279
+ request.url.path,
280
+ str(e),
281
+ )
282
+ raise
283
+
284
+
285
+###############################################################################
286
+# FastAPI lifecycle
287
+###############################################################################
288
+
289
+
290
+@app.on_event("startup")
291
+async def on_startup() -> None:
292
+ global encode_session, prefill_session, decode_session
293
+ timeout = aiohttp.ClientTimeout(total=100_000)
294
+ connector = aiohttp.TCPConnector(limit=0, force_close=False)
295
+ encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
296
+ if app.state.p_urls:
297
+ # only setup if prefill instance(s) exist
298
+ prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
299
+ decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
300
+
301
+
302
+@app.on_event("shutdown")
303
+async def on_shutdown() -> None:
304
+ global encode_session, prefill_session, decode_session
305
+ if encode_session:
306
+ await encode_session.close()
307
+ if prefill_session:
308
+ await prefill_session.close()
309
+ if decode_session:
310
+ await decode_session.close()
311
+
312
+
313
+###############################################################################
314
+# Core forwarding
315
+###############################################################################
316
+
317
+
318
+async def forward_non_stream(
319
+ req_data: dict, req_id: str, e_urls: list[str], p_url: str, d_url: str
320
+) -> dict:
321
+ try:
322
+ # Step 1: Process through Encoder instance (if has MM input)
323
+ await fanout_encoder_primer(req_data, e_urls, req_id)
324
+
325
+ # Step 2: Process through Prefill instance
326
+ req_data = await maybe_prefill(req_data, p_url, req_id)
327
+
328
+ # Step 3: Process through Decode instance
329
+ logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
330
+ headers = {"x-request-id": req_id}
331
+
332
+ # Non-streaming response
333
+ async with decode_session.post(
334
+ f"{d_url}/v1/chat/completions", json=req_data, headers=headers
335
+ ) as resp:
336
+ resp.raise_for_status()
337
+ return await resp.json()
338
+
339
+ except HTTPException:
340
+ raise
341
+ except Exception as e:
342
+ logger.exception("[%s] Error in forward_non_stream: %s", req_id, str(e))
343
+ raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") from e
344
+
345
+
346
+async def forward_stream(
347
+ req_data: dict, req_id: str, e_urls: list[str], p_url: str, d_url: str
348
+) -> AsyncIterator[str]:
349
+ try:
350
+ # Step 1: Process through Encoder instance (if has MM input)
351
+ await fanout_encoder_primer(req_data, e_urls, req_id)
352
+
353
+ # Step 2: Process through Prefill instance
354
+ req_data = await maybe_prefill(req_data, p_url, req_id)
355
+
356
+ # Step 3: Process through Decode instance
357
+ logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
358
+ headers = {"x-request-id": req_id}
359
+
360
+ # Streaming response
361
+ async with decode_session.post(
362
+ f"{d_url}/v1/chat/completions",
363
+ json=req_data,
364
+ headers=headers,
365
+ ) as resp:
366
+ resp.raise_for_status()
367
+ async for chunk in resp.content.iter_chunked(1024):
368
+ if chunk:
369
+ yield chunk.decode("utf-8", errors="ignore")
370
+
371
+ logger.info("[%s] Streaming completed", req_id)
372
+
373
+ except HTTPException:
374
+ logger.exception("[%s] HTTPException in forward_stream", req_id)
375
+ raise
376
+ except Exception as e:
377
+ logger.exception("[%s] Error in forward_stream: %s", req_id, str(e))
378
+ raise HTTPException(
379
+ status_code=500, detail=f"Proxy streaming error: {str(e)}"
380
+ ) from e
381
+
382
+
383
+###############################################################################
384
+# Public routes
385
+###############################################################################
386
+
387
+
388
+@app.post("/v1/chat/completions")
389
+async def chat_completions(request: Request):
390
+ try:
391
+ req_data = await request.json()
392
+ req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
393
+
394
+ e_urls = app.state.e_urls # we want the full list for fan-out
395
+ p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
396
+ d_url = random.choice(app.state.d_urls)
397
+
398
+ is_streaming = req_data.get("stream", False)
399
+
400
+ if is_streaming:
401
+ return StreamingResponse(
402
+ forward_stream(req_data, req_id, e_urls, p_url, d_url),
403
+ media_type="text/event-stream",
404
+ )
405
+ result = await forward_non_stream(req_data, req_id, e_urls, p_url, d_url)
406
+ return JSONResponse(content=result)
407
+
408
+ except HTTPException:
409
+ raise
410
+ except Exception as e:
411
+ logger.exception("Error in chat_completions endpoint: %s", str(e))
412
+ raise HTTPException(
413
+ status_code=500, detail=f"Request processing error: {str(e)}"
414
+ ) from e
415
+
416
+
417
+@app.get("/v1/models")
418
+async def list_models():
419
+ async with decode_session.get(f"{app.state.d_urls[0]}/v1/models") as resp:
420
+ resp.raise_for_status()
421
+ return await resp.json()
422
+
423
+
424
+@app.get("/health")
425
+async def health_check():
426
+ async def healthy(urls):
427
+ if not urls:
428
+ return "empty"
429
+ for u in urls:
430
+ try:
431
+ async with encode_session.get(f"{u}/health") as resp:
432
+ resp.raise_for_status()
433
+ except Exception:
434
+ return "unhealthy"
435
+ return "healthy"
436
+
437
+ e_status, p_status, d_status = await asyncio.gather(
438
+ healthy(app.state.e_urls), healthy(app.state.p_urls), healthy(app.state.d_urls)
439
+ )
440
+
441
+ overall_healthy = all(
442
+ status != "unhealthy" for status in (e_status, p_status, d_status)
443
+ )
444
+
445
+ status_code = 200 if overall_healthy else 503
446
+
447
+ return JSONResponse(
448
+ {
449
+ "proxy": "healthy",
450
+ "encode_cluster": e_status,
451
+ "prefill_cluster": p_status,
452
+ "decode_cluster": d_status,
453
+ },
454
+ status_code=status_code,
455
+ )
456
+
457
+
458
+###############################################################################
459
+# Simple profiler fan-out (unchanged except for sessions)
460
+###############################################################################
461
+
462
+
463
+async def _post_if_available(
464
+ session: aiohttp.ClientSession,
465
+ url: str,
466
+ payload: dict,
467
+ headers: dict,
468
+) -> dict | None:
469
+ """
470
+ POST `payload` to `url`.
471
+
472
+ Returns
473
+ -------
474
+ • The decoded JSON body on success (2xx)
475
+ • None if the endpoint does not exist (404)
476
+ • Raises for anything else.
477
+ """
478
+ try:
479
+ resp = await session.post(url, json=payload, headers=headers)
480
+ if resp.status == 404: # profiling disabled on that server
481
+ logger.warning("Profiling endpoint missing on %s", url)
482
+ return None
483
+ resp.raise_for_status()
484
+ return await resp.json(content_type=None)
485
+ except aiohttp.ClientResponseError as exc:
486
+ # Pass 404 through the branch above, re-raise everything else
487
+ if exc.status == 404:
488
+ logger.warning("Profiling endpoint missing on %s", url)
489
+ return None
490
+ raise
491
+ except Exception:
492
+ # Network errors etc.: propagate
493
+ raise
494
+
495
+
496
+async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str):
497
+ """
498
+ Fire & forget to both clusters, tolerate 404.
499
+ """
500
+ headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
501
+
502
+ encode_task = _post_if_available(
503
+ encode_session, f"{e_url}/{cmd}_profile", payload, headers
504
+ )
505
+ prefill_task = (
506
+ _post_if_available(prefill_session, f"{p_url}/{cmd}_profile", payload, headers)
507
+ if p_url is not None
508
+ else asyncio.sleep(0)
509
+ )
510
+ decode_task = _post_if_available(
511
+ decode_session, f"{d_url}/{cmd}_profile", payload, headers
512
+ )
513
+
514
+ encode_res, prefill_res, decode_res = await asyncio.gather(
515
+ encode_task, prefill_task, decode_task
516
+ )
517
+
518
+ # If *all* clusters said “I don’t have that route”, surface an error
519
+ if encode_res is prefill_res is decode_res is None:
520
+ raise HTTPException(
521
+ status_code=503,
522
+ detail="Profiling endpoints are disabled on all clusters",
523
+ )
524
+
525
+ return {
526
+ "encode": encode_res, # may be None
527
+ "prefill": prefill_res, # may be None
528
+ "decode": decode_res, # may be None
529
+ }
530
+
531
+
532
+@app.post("/start_profile")
533
+async def start_profile(request: Request):
534
+ body = await request.json()
535
+ # TODO: handle multi urls properly
536
+ e_url = random.choice(app.state.e_urls)
537
+ p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
538
+ d_url = random.choice(app.state.d_urls)
539
+ return await _profile_cmd("start", body, e_url, p_url, d_url)
540
+
541
+
542
+@app.post("/stop_profile")
543
+async def stop_profile(request: Request):
544
+ body = await request.json()
545
+ # TODO: handle multi urls properly
546
+ e_url = random.choice(app.state.e_urls)
547
+ p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
548
+ d_url = random.choice(app.state.d_urls)
549
+ return await _profile_cmd("stop", body, e_url, p_url, d_url)
550
+
551
+
552
+if __name__ == "__main__":
553
+ parser = argparse.ArgumentParser()
554
+ parser.add_argument("--host", default="0.0.0.0")
555
+ parser.add_argument("--port", type=int, default=8000)
556
+ parser.add_argument(
557
+ "--encode-servers-urls",
558
+ required=True,
559
+ help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")',
560
+ )
561
+ parser.add_argument(
562
+ "--prefill-servers-urls",
563
+ required=True,
564
+ help=(
565
+ 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ',
566
+ 'to enable E->P->D, set "disable" or "none" to enable E->PD',
567
+ ),
568
+ )
569
+ parser.add_argument(
570
+ "--decode-servers-urls",
571
+ required=True,
572
+ help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
573
+ )
574
+
575
+ args = parser.parse_args()
576
+ app.state.e_urls = [
577
+ u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
578
+ ]
579
+ app.state.d_urls = [
580
+ u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
581
+ ]
582
+ # handle prefill instances
583
+ if args.prefill_servers_urls.lower() in ("disable", "none", ""):
584
+ app.state.p_urls = []
585
+ logger.info(
586
+ "Disaggregated prefill phase explicitly disabled by user. Running E + PD..."
587
+ )
588
+ else:
589
+ app.state.p_urls = [
590
+ u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
591
+ ]
592
+ logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
593
+
594
+ logger.info("Proxy listening on %s:%s", args.host, args.port)
595
+ logger.info("Encode servers: %s", app.state.e_urls)
596
+ logger.info("Prefill instances %s", app.state.p_urls)
597
+ logger.info("Decode servers: %s", app.state.d_urls)
598
+
599
+ uvicorn.run(
600
+ app,
601
+ host=args.host,
602
+ port=args.port,
603
+ log_level="info",
604
+ loop="uvloop",
605
+ access_log=True,
606
+ )
tests/v1/core/test_scheduler.py
+1083
-16
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.config import (
10
CacheConfig,
11
+ ECTransferConfig,
12
KVTransferConfig,
13
ModelConfig,
14
SchedulerConfig,
@@ -20,6 +21,9 @@ from vllm.multimodal.inputs import (
21
PlaceholderRange,
22
)
23
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
24
+from vllm.utils.hashing import sha256
25
+from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
26
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
27
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
28
from vllm.v1.core.sched.scheduler import Scheduler
29
from vllm.v1.kv_cache_interface import (
@@ -872,7 +876,10 @@ def _step_until_done(
876
for _, num_scheduled_tokens in output.num_scheduled_tokens.items():
877
# We should be in the decode phase now.
878
assert num_scheduled_tokens == 1
875
- assert len(output.kv_connector_metadata.requests) == 0
879
+ if scheduler.connector is not None:
880
+ assert len(output.kv_connector_metadata.requests) == 0
881
+ if scheduler.ec_connector is not None:
882
+ assert len(output.ec_connector_metadata.mm_datas) == 0
883
ecos = scheduler.update_from_output(output, model_runner_output)[0]
884
all_done = True
885
for eco in ecos.outputs:
@@ -1066,7 +1073,10 @@ def test_external_prefix_cache_metrics():
1073
assert external_stats.preempted_requests == 0
1074
1075
1069
-def test_kv_connector_unable_to_allocate():
1076
+@pytest.mark.parametrize(
1077
+ "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
1078
+)
1079
+def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
1080
"""
1081
Test whether scheduler with KVConnector is able to handle
1082
unable to allocate (run out of blocks in allocate_slots().
@@ -1080,6 +1090,9 @@ def test_kv_connector_unable_to_allocate():
1090
use_kv_connector=True,
1091
block_size=BLOCK_SIZE,
1092
num_blocks=NUM_BLOCKS,
1093
+ # encoder connector should not affect test results
1094
+ use_ec_connector=use_ec_connector,
1095
+ ec_role=ec_role,
1096
)
1097
NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
1098
scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
@@ -1148,7 +1161,10 @@ def test_kv_connector_unable_to_allocate():
1161
assert len(scheduler.waiting) == 0
1162
1163
1151
-def test_kv_connector_handles_preemption():
1164
+@pytest.mark.parametrize(
1165
+ "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
1166
+)
1167
+def test_kv_connector_handles_preemption(use_ec_connector, ec_role):
1168
"""
1169
Test whether scheduler with KVConnector is able to handle
1170
unable to allocate (run out of blocks in allocate_slots().
@@ -1163,6 +1179,9 @@ def test_kv_connector_handles_preemption():
1179
use_kv_connector=True,
1180
block_size=BLOCK_SIZE,
1181
num_blocks=NUM_BLOCKS,
1182
+ # encoder connector should not affect test results
1183
+ use_ec_connector=use_ec_connector,
1184
+ ec_role=ec_role,
1185
)
1186
1187
NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
@@ -1379,6 +1398,8 @@ def create_scheduler_with_priority(
1398
block_size: int = 16,
1399
max_model_len: int | None = None,
1400
num_speculative_tokens: int | None = None,
1401
+ use_ec_connector: bool = False,
1402
+ ec_role: str | None = None,
1403
) -> Scheduler:
1404
"""Create scheduler with priority policy enabled.
1405
@@ -1439,12 +1460,23 @@ def create_scheduler_with_priority(
1460
model="ngram", num_speculative_tokens=num_speculative_tokens
1461
)
1462
1463
+ ec_transfer_config = (
1464
+ ECTransferConfig(
1465
+ ec_connector="ECSharedStorageConnector",
1466
+ ec_role=ec_role,
1467
+ ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test"},
1468
+ )
1469
+ if use_ec_connector
1470
+ else None
1471
+ )
1472
+
1473
vllm_config = VllmConfig(
1474
scheduler_config=scheduler_config,
1475
model_config=model_config,
1476
cache_config=cache_config,
1477
kv_transfer_config=kv_transfer_config,
1478
speculative_config=speculative_config,
1479
+ ec_transfer_config=ec_transfer_config,
1480
)
1481
kv_cache_config = KVCacheConfig(
1482
num_blocks=num_blocks, # A large number of blocks to hold all requests
@@ -1465,16 +1497,23 @@ def create_scheduler_with_priority(
1497
)
1498
1499
1500
+_none_hash_initialized = False
1501
+
1502
+
1503
def create_requests_with_priority(
1504
num_requests: int,
1505
priorities: list[int],
1506
arrival_times: list[float] | None = None,
1507
num_tokens: int = 10,
1508
+ mm_hashes_list: list[list[str]] | None = None,
1509
mm_positions: list[list[PlaceholderRange]] | None = None,
1510
max_tokens: int = 16,
1511
stop_token_ids: list[int] | None = None,
1512
prompt_logprobs: int | None = None,
1513
starting_idx: int = 0,
1514
+ same_prompt: bool = False,
1515
+ block_size: int = 16,
1516
+ req_ids: list[str] | None = None,
1517
):
1518
"""Create requests with specified priorities and arrival times."""
1519
assert len(priorities) == num_requests
@@ -1483,6 +1522,12 @@ def create_requests_with_priority(
1522
else:
1523
arrival_times = [float(i) for i in range(num_requests)]
1524
1525
+ global _none_hash_initialized
1526
+ if not _none_hash_initialized:
1527
+ init_none_hash(sha256)
1528
+ _none_hash_initialized = True
1529
+
1530
+ block_hasher = get_request_block_hasher(block_size, sha256)
1531
sampling_params = SamplingParams(
1532
ignore_eos=False,
1533
max_tokens=max_tokens,
@@ -1490,29 +1535,70 @@ def create_requests_with_priority(
1535
prompt_logprobs=prompt_logprobs,
1536
)
1537
requests = []
1538
+
1539
+ if mm_hashes_list is not None:
1540
+ # NOTE: allow manual input; some mm items can have the same identifier
1541
+ # no. of mm_hashes and mm_positions for each request should be identical
1542
+ assert mm_positions is not None, (
1543
+ "mm_positions must be provided when mm_hashes_list is provided"
1544
+ )
1545
+ assert len(mm_hashes_list) == len(mm_positions) == num_requests
1546
+ assert [len(h) for h in mm_hashes_list] == [len(p) for p in mm_positions]
1547
+
1548
+ # Since same identifier would imply they are identical encoder output
1549
+ # Verify mm items with identical identifier are having mm_position.length
1550
+ seen_hashes: dict[str, int] = {}
1551
+
1552
+ if req_ids:
1553
+ assert len(req_ids) == num_requests
1554
+ else:
1555
+ req_ids = [f"{i + starting_idx}" for i in range(num_requests)]
1556
+
1557
for i in range(num_requests):
1558
mm_features = []
1495
- if mm_positions is not None:
1496
- mm_position = mm_positions[i]
1497
- for j, position in enumerate(mm_position):
1559
+
1560
+ for j, position in enumerate(
1561
+ mm_positions[i] if mm_positions is not None else []
1562
+ ):
1563
+ if mm_hashes_list is not None:
1564
+ identifier = mm_hashes_list[i][j]
1565
+
1566
+ # Verify if position length is identical
1567
+ position_length = position.length
1568
+ if identifier in seen_hashes:
1569
+ assert seen_hashes[identifier] == position_length, (
1570
+ f"mm_hash '{identifier}' has inconsistent position lengths: "
1571
+ f"previously {seen_hashes[identifier]}, now {position_length} "
1572
+ f"at request {i}, position {j}"
1573
+ )
1574
+ else:
1575
+ seen_hashes[identifier] = position_length
1576
+ else:
1577
+ # Unique dummy hash for each mm item
1578
identifier = f"hash{i}_{j}"
1499
- mm_feature = MultiModalFeatureSpec(
1500
- data=MultiModalKwargsItem.dummy("dummy_m"),
1501
- mm_position=position,
1502
- identifier=identifier,
1503
- modality="image",
1504
- )
1505
- mm_features.append(mm_feature)
1579
+ mm_feature = MultiModalFeatureSpec(
1580
+ data=MultiModalKwargsItem.dummy("dummy_m"),
1581
+ mm_position=position,
1582
+ identifier=identifier,
1583
+ modality="image",
1584
+ )
1585
+ mm_features.append(mm_feature)
1586
1587
+ prompt_token_ids = (
1588
+ [starting_idx] * num_tokens
1589
+ if same_prompt
1590
+ else [i + starting_idx] * num_tokens
1591
+ )
1592
request = Request(
1508
- request_id=f"{i + starting_idx}",
1509
- prompt_token_ids=[i + starting_idx] * num_tokens,
1593
+ request_id=req_ids[i],
1594
+ prompt_token_ids=prompt_token_ids,
1595
sampling_params=sampling_params,
1596
pooling_params=None,
1597
mm_features=mm_features if mm_features else None,
1598
eos_token_id=EOS_TOKEN_ID,
1599
arrival_time=arrival_times[i],
1600
priority=priorities[i],
1601
+ block_hasher=block_hasher,
1602
)
1603
requests.append(request)
1604
return requests
@@ -1999,7 +2085,12 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
2085
assert len(scheduler.waiting) == 1
2086
2087
2002
-def test_priority_scheduling_preemption_and_resumption_when_out_of_kv():
2088
+@pytest.mark.parametrize(
2089
+ "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
2090
+)
2091
+def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
2092
+ use_ec_connector, ec_role
2093
+):
2094
"""Test that priority scheduling preempts lower priority requests
2095
when out of KV cache space."""
2096
# Create scheduler with very limited memory to force preemption
@@ -2009,6 +2100,9 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv():
2100
num_blocks=5, # Can hold 64 tokens (first block is null)
2101
block_size=16, # Standard block size
2102
use_kv_connector=True,
2103
+ # encoder connector should not affect test results
2104
+ use_ec_connector=use_ec_connector,
2105
+ ec_role=ec_role,
2106
)
2107
2108
# Create a request and schedule it
@@ -2168,3 +2262,976 @@ def _validate_chunked_prefill_settings_for_encoder_decoder(
2262
assert scheduler_config.disable_chunked_mm_input is not expect_enabled
2263
if is_encoder_decoder and not expect_enabled:
2264
assert scheduler_config.long_prefill_token_threshold == 0
2265
+
2266
+
2267
+# ==============================================================================
2268
+# EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests start
2269
+# NOTE: In E->P->D disagg case, both KV and EC Connector works in P instance
2270
+# Unless specify, the existence of KV Connector should not affect any test results
2271
+# ==============================================================================
2272
+
2273
+
2274
+def _assert_right_encoder_cache_allocated(
2275
+ scheduler: Scheduler,
2276
+ hashes_to_check: list[str] | None = None,
2277
+ requests: list[Request] | None = None,
2278
+ expected_total_allocated: int | None = None,
2279
+):
2280
+ """Check whether encoder cache is allocated correctly."""
2281
+ encoder_cache_manager = scheduler.encoder_cache_manager
2282
+
2283
+ # Verify encoder cache manager exists
2284
+ assert encoder_cache_manager is not None, "Encoder cache manager should exist"
2285
+
2286
+ # Verify number of cache
2287
+ if expected_total_allocated is not None:
2288
+ assert len(encoder_cache_manager.cached) == expected_total_allocated
2289
+ if expected_total_allocated == 0:
2290
+ return
2291
+
2292
+ # Verify each request with MM data is in cache
2293
+ cached_hashes = set(encoder_cache_manager.cached.keys())
2294
+
2295
+ if hashes_to_check:
2296
+ missed_hashes = set(hashes_to_check) - cached_hashes
2297
+ assert not missed_hashes, (
2298
+ f"Miss hashes: {missed_hashes} "
2299
+ f"Existing encoder cache: {encoder_cache_manager.cached}"
2300
+ )
2301
+
2302
+ for req in requests if requests is not None else []:
2303
+ if req.mm_features:
2304
+ mm_hashes = [f.identifier for f in req.mm_features]
2305
+ req_hashes = set(mm_hashes) # unique hashes set
2306
+ missed_hashes = req_hashes - cached_hashes
2307
+ assert not missed_hashes, (
2308
+ f"Miss hashes in cache for request {req.request_id}: {missed_hashes} "
2309
+ f"Existing encoder cache: {encoder_cache_manager.cached}"
2310
+ )
2311
+
2312
+
2313
+def _assert_right_ec_connector_metadata(
2314
+ output: SchedulerOutput,
2315
+ mm_features_list: list[MultiModalFeatureSpec],
2316
+):
2317
+ """Verify that ECConnector metadata EXACTLY matches the input MM data"""
2318
+ # Get the connector metadata
2319
+ metadata = output.ec_connector_metadata
2320
+
2321
+ # Create lookup dictionaries for efficient access
2322
+ metadata_dict = {mm_data.mm_hash: mm_data for mm_data in metadata.mm_datas}
2323
+
2324
+ # Check all required identifiers exist in metadata; and no extra
2325
+ # In ECSharedStorageConnector format
2326
+ # NOTE: even having same identifier, the mm_features can be different
2327
+ # since their mm_position can be in different offsets, etc
2328
+ identifiers_dict = {f.identifier for f in mm_features_list}
2329
+ assert set(metadata_dict.keys()) == identifiers_dict
2330
+
2331
+ # Verify the info matches
2332
+ for i, mm_feature in enumerate(mm_features_list):
2333
+ identifier = mm_feature.identifier
2334
+ assert metadata_dict[identifier].mm_hash == identifier
2335
+ assert metadata_dict[identifier].num_token == mm_feature.mm_position.length
2336
+
2337
+
2338
+def _assert_right_encoder_inputs(
2339
+ output: SchedulerOutput,
2340
+ check_exist: bool | None = True,
2341
+ requests: list[Request] | None = None,
2342
+ expected_encoder_inputs: list[list[int]] | None = None,
2343
+ expected_total_reqs: int | None = None,
2344
+):
2345
+ """Verify that requests/mm_hashes should (not) in scheduled encoder input
2346
+ If check_exist is False, this function returns True
2347
+ if requests are NOT in encoder inputs"""
2348
+
2349
+ # Get the scheduled encoder inputs
2350
+ # NOTE: scheduled_encoder_inputs is a dictionary with request id as key
2351
+ scheduled_encoder_inputs = output.scheduled_encoder_inputs
2352
+
2353
+ # Check if scheduled_encoder_inputs is empty as expected
2354
+ if expected_total_reqs is not None:
2355
+ assert len(scheduled_encoder_inputs) == expected_total_reqs
2356
+ if expected_total_reqs == 0:
2357
+ return
2358
+
2359
+ # Number of expected enocder inputs should match number of requests
2360
+ if expected_encoder_inputs:
2361
+ assert check_exist and requests is not None # only support expect input exist
2362
+ assert len(requests) == len(expected_encoder_inputs)
2363
+
2364
+ # Check request (not) exist as expected
2365
+ for i, request in enumerate(requests if requests is not None else []):
2366
+ assert (request.request_id in scheduled_encoder_inputs) is check_exist, (
2367
+ f"Request {request.id} presence mismatch: expected {check_exist}, "
2368
+ f"got {request.id in scheduled_encoder_inputs}"
2369
+ )
2370
+ if expected_encoder_inputs:
2371
+ scheduled_encoder_input = scheduled_encoder_inputs[request.request_id]
2372
+ assert scheduled_encoder_input == expected_encoder_inputs[i]
2373
+
2374
+
2375
+def test_scheduler_no_ec_connector_by_default():
2376
+ """Test scheduler doesn't have EC connector by default."""
2377
+ scheduler = create_scheduler()
2378
+ assert scheduler.ec_connector is None
2379
+
2380
+
2381
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2382
+def test_ec_connector_text_only_request(use_kv_connector):
2383
+ """Test text-only requests don't allocate encoder cache."""
2384
+ scheduler = create_scheduler(
2385
+ model="llava-hf/llava-1.5-7b-hf",
2386
+ use_kv_connector=use_kv_connector,
2387
+ use_ec_connector=True,
2388
+ ec_role="ec_consumer",
2389
+ )
2390
+
2391
+ NUM_PROMPT_TOKENS = 100
2392
+
2393
+ # Create text-only request (no mm_positions)
2394
+ requests = create_requests(
2395
+ num_requests=1,
2396
+ num_tokens=NUM_PROMPT_TOKENS,
2397
+ )
2398
+ assert not requests[0].mm_features # No MM data
2399
+
2400
+ scheduler.add_request(requests[0])
2401
+ output = scheduler.schedule()
2402
+
2403
+ # Should schedule
2404
+ assert len(output.scheduled_new_reqs) == 1
2405
+
2406
+ # Scheduled tokens should equal prompt tokens exactly
2407
+ scheduled = output.num_scheduled_tokens[requests[0].request_id]
2408
+ assert scheduled == NUM_PROMPT_TOKENS, (
2409
+ f"Text-only should schedule {NUM_PROMPT_TOKENS}, got {scheduled}"
2410
+ )
2411
+
2412
+ # Encoder cache should be empty
2413
+ _assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0)
2414
+
2415
+ # ECConnector should carry no metadata
2416
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
2417
+
2418
+ # Scheduled encoder input should be empty; no mm to compute
2419
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2420
+
2421
+
2422
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2423
+def test_ec_connector_cache_hit_external_load(use_kv_connector):
2424
+ """Test ec_consumer loads from external cache when hit.
2425
+ A normal basic operation for EPD disaggrgation"""
2426
+ scheduler = create_scheduler(
2427
+ model="llava-hf/llava-1.5-7b-hf",
2428
+ enable_prefix_caching=True,
2429
+ # kv connector should not effect test results
2430
+ use_kv_connector=use_kv_connector,
2431
+ use_ec_connector=True,
2432
+ ec_role="ec_consumer",
2433
+ )
2434
+
2435
+ # Create MM request
2436
+ NUM_TOKENS = 200 # NOTE: includes mm tokens
2437
+ NUM_ENCODER_TOKENS = 100
2438
+ mm_hashes_list = [["hash_test1"]]
2439
+ mm_positions = [[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS)]]
2440
+
2441
+ request = create_requests(
2442
+ num_requests=1,
2443
+ num_tokens=NUM_TOKENS,
2444
+ mm_hashes_list=mm_hashes_list,
2445
+ mm_positions=mm_positions,
2446
+ )[0]
2447
+
2448
+ # Mock cache hit - encoder cache exists externally
2449
+ scheduler.ec_connector.has_caches = Mock(return_value=[True])
2450
+ scheduler.ec_connector.update_state_after_alloc = Mock(
2451
+ wraps=scheduler.ec_connector.update_state_after_alloc
2452
+ )
2453
+
2454
+ scheduler.add_request(request)
2455
+ output = scheduler.schedule()
2456
+
2457
+ # Should schedule prompt tokens
2458
+ scheduled_tokens = output.num_scheduled_tokens[request.request_id]
2459
+ assert scheduled_tokens == NUM_TOKENS
2460
+
2461
+ # Should called update_state_after_alloc for external load
2462
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request, 0)
2463
+
2464
+ # Encoder cache should contain mm items from request
2465
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request])
2466
+
2467
+ # ECConnector should carry metadata of request
2468
+ _assert_right_ec_connector_metadata(output, mm_features_list=request.mm_features)
2469
+
2470
+ # Scheduled encoder input should be empty; no mm to compute
2471
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2472
+
2473
+
2474
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2475
+def test_ec_connector_cache_miss_computes_locally(use_kv_connector):
2476
+ """Test consumer can compute encoder locally when cache miss (fallback)."""
2477
+ # encoder cache itself if it doesn't receive it from external storage
2478
+
2479
+ scheduler = create_scheduler(
2480
+ model="llava-hf/llava-1.5-7b-hf",
2481
+ enable_prefix_caching=True,
2482
+ use_kv_connector=use_kv_connector,
2483
+ use_ec_connector=True,
2484
+ ec_role="ec_consumer",
2485
+ )
2486
+
2487
+ # Verify consumer role
2488
+ assert scheduler.ec_connector is not None
2489
+ assert not scheduler.ec_connector.is_producer
2490
+
2491
+ # Create MM request
2492
+ request_mm_missed = create_requests(
2493
+ num_requests=1,
2494
+ num_tokens=200, # Total (including 100 MM)
2495
+ mm_positions=[[PlaceholderRange(offset=0, length=100)]], # 100 MM tokens
2496
+ )[0]
2497
+
2498
+ # Mock cache miss - encoder cache doesn't exist externally
2499
+ scheduler.ec_connector.has_caches = Mock(return_value=[False])
2500
+
2501
+ scheduler.add_request(request_mm_missed)
2502
+ output = scheduler.schedule()
2503
+
2504
+ # SCHEDULER should decide to compute encoder locally (fallback)
2505
+ assert len(output.scheduled_new_reqs) == 1
2506
+
2507
+ # Should schedule full prompt tokens
2508
+ scheduled_tokens = output.num_scheduled_tokens[request_mm_missed.request_id]
2509
+ assert scheduled_tokens == 200, (
2510
+ f"Expected 200 tokens on cache miss, got {scheduled_tokens}"
2511
+ )
2512
+
2513
+ # Encoder cache should contain mm items from request
2514
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request_mm_missed])
2515
+
2516
+ # ECConnector should carry no metadata (missed cache)
2517
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
2518
+
2519
+ # Scheduled encoder input contain mm for request_mm_missed
2520
+ _assert_right_encoder_inputs(
2521
+ output,
2522
+ requests=[request_mm_missed],
2523
+ expected_encoder_inputs=[[0]], # index 0 of the mm item
2524
+ expected_total_reqs=1,
2525
+ )
2526
+
2527
+ # Then MODEL_RUNNER will execute the encoder and cache the result
2528
+
2529
+
2530
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2531
+def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):
2532
+ """Test consumer with partial cache hit (local & connector) with 2 requests."""
2533
+ scheduler = create_scheduler(
2534
+ model="llava-hf/llava-1.5-7b-hf",
2535
+ enable_prefix_caching=True,
2536
+ use_kv_connector=use_kv_connector,
2537
+ use_ec_connector=True,
2538
+ ec_role="ec_consumer",
2539
+ )
2540
+
2541
+ # Create MM request
2542
+ NUM_TOKENS_1 = 300 # NOTE: includes mm tokens
2543
+ NUM_ENCODER_TOKENS_1 = 50
2544
+ mm_hashes_list_1 = [["hash1_A", "hash1_B", "hash1_A", "hash1_F"]]
2545
+ mm_positions_1 = [
2546
+ [
2547
+ PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_1),
2548
+ PlaceholderRange(offset=100, length=NUM_ENCODER_TOKENS_1),
2549
+ PlaceholderRange(offset=200, length=NUM_ENCODER_TOKENS_1),
2550
+ PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_1),
2551
+ ]
2552
+ ]
2553
+
2554
+ # Create request with 4 MM items, with 2 identical items
2555
+ request1 = create_requests(
2556
+ num_requests=1,
2557
+ num_tokens=NUM_TOKENS_1,
2558
+ mm_hashes_list=mm_hashes_list_1,
2559
+ mm_positions=mm_positions_1,
2560
+ max_tokens=1, # For simplicity
2561
+ )[0]
2562
+
2563
+ # Mock partial cache hit: 1st and 3rd missing, 2nd and 4th exist
2564
+ scheduler.ec_connector.has_caches = Mock(return_value=[False, True, False, True])
2565
+ scheduler.ec_connector.update_state_after_alloc = Mock(
2566
+ wraps=scheduler.ec_connector.update_state_after_alloc
2567
+ )
2568
+
2569
+ scheduler.add_request(request1)
2570
+ output = scheduler.schedule()
2571
+
2572
+ # Should schedule all tokens
2573
+ scheduled_tokens = output.num_scheduled_tokens[request1.request_id]
2574
+ assert scheduled_tokens == NUM_TOKENS_1
2575
+
2576
+ # Encoder cache should contain all mm items from request
2577
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request1])
2578
+
2579
+ # Should have called update_state_after_alloc for external load
2580
+ scheduler.ec_connector.update_state_after_alloc.assert_called()
2581
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2582
+
2583
+ # ECConnector should carry metadata for 2nd and 4th mm item
2584
+ _assert_right_ec_connector_metadata(
2585
+ output, mm_features_list=[request1.mm_features[1], request1.mm_features[3]]
2586
+ )
2587
+
2588
+ # Should schedule ONLY 1 encoder input (index 0), no repeat for identical items
2589
+ _assert_right_encoder_inputs(
2590
+ output,
2591
+ requests=[request1],
2592
+ expected_encoder_inputs=[[0]], # index 0 of the mm item ONLY
2593
+ expected_total_reqs=1,
2594
+ )
2595
+
2596
+ # Simulate model execution 1 step
2597
+ model_output = ModelRunnerOutput(
2598
+ req_ids=[request1.request_id],
2599
+ req_id_to_index={request1.request_id: 0},
2600
+ sampled_token_ids=[[100]],
2601
+ # spec_token_ids=None,
2602
+ logprobs=None,
2603
+ prompt_logprobs_dict={},
2604
+ pooler_output=[],
2605
+ )
2606
+ scheduler.update_from_output(output, model_output)
2607
+
2608
+ # request1 is finished after outputing 1 token
2609
+ # Finish request
2610
+ scheduler.finish_requests(request1.request_id, RequestStatus.FINISHED_LENGTH_CAPPED)
2611
+
2612
+ # Create another request with 4 MM items
2613
+ NUM_TOKENS_2 = 400
2614
+ NUM_ENCODER_TOKENS_2 = 50
2615
+ mm_hashes_list_2 = [["hash1_C", "hash1_D", "hash1_E", "hash1_A"]]
2616
+ mm_positions_2 = [
2617
+ [
2618
+ PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_2),
2619
+ PlaceholderRange(offset=100, length=NUM_ENCODER_TOKENS_2),
2620
+ PlaceholderRange(offset=200, length=NUM_ENCODER_TOKENS_2),
2621
+ PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_2),
2622
+ ]
2623
+ ]
2624
+
2625
+ request2 = create_requests(
2626
+ num_requests=1,
2627
+ num_tokens=NUM_TOKENS_2,
2628
+ mm_hashes_list=mm_hashes_list_2,
2629
+ mm_positions=mm_positions_2,
2630
+ max_tokens=1, # For simplicity
2631
+ )[0]
2632
+
2633
+ # Mock partial cache hit: only hash1_A and hash1_C exist in connector
2634
+ scheduler.ec_connector.has_caches = Mock(return_value=[True, False, False, True])
2635
+
2636
+ scheduler.add_request(request2)
2637
+ output = scheduler.schedule()
2638
+
2639
+ # Check
2640
+ # Should schedule all tokens
2641
+ scheduled_tokens = output.num_scheduled_tokens[request2.request_id]
2642
+ assert scheduled_tokens == 400
2643
+
2644
+ # Encoder cache should contain all mm items from request2
2645
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request2])
2646
+
2647
+ # Should call update_state_after_alloc for hash1_C, ONLY
2648
+ # hash1_A should not be loaded from connector
2649
+ # since it's computed in last request & exist in local cache
2650
+ # Order of getting encoder cache should be: local cache -> connector-> compute
2651
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request2, 0)
2652
+ scheduler.ec_connector.update_state_after_alloc.assert_called_once()
2653
+
2654
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2655
+
2656
+ # ECConnector should carry metadata for hash1_C only (index 0)
2657
+ _assert_right_ec_connector_metadata(
2658
+ output, mm_features_list=[request2.mm_features[0]]
2659
+ )
2660
+
2661
+ # Should schedule 2 encoder input hash1_D and hash1_E (index 1, 2)
2662
+ _assert_right_encoder_inputs(
2663
+ output,
2664
+ requests=[request2],
2665
+ expected_encoder_inputs=[[1, 2]],
2666
+ expected_total_reqs=1,
2667
+ )
2668
+
2669
+
2670
+@pytest.mark.parametrize("cache_exist", ["local", "connector_only", "no_where"])
2671
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2672
+def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):
2673
+ scheduler = create_scheduler(
2674
+ model="llava-hf/llava-1.5-7b-hf",
2675
+ max_num_seqs=10, # allow multiple requests
2676
+ max_num_batched_tokens=2048,
2677
+ enable_prefix_caching=True,
2678
+ use_kv_connector=use_kv_connector,
2679
+ use_ec_connector=True,
2680
+ ec_role="ec_consumer",
2681
+ )
2682
+ mm_hashes_list = [[f"hash_{i}"] for i in range(10)]
2683
+ mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)]
2684
+ requests = create_requests(
2685
+ num_requests=10,
2686
+ num_tokens=200,
2687
+ mm_hashes_list=mm_hashes_list,
2688
+ mm_positions=mm_positions,
2689
+ )
2690
+ for request in requests:
2691
+ scheduler.add_request(request)
2692
+
2693
+ # Set up to test different encoder cache exsistence scenario after preemption
2694
+ # Order of getting encoder cache should be: local cache -> connector-> compute
2695
+ scheduler.ec_connector.update_state_after_alloc = Mock(
2696
+ wraps=scheduler.ec_connector.update_state_after_alloc
2697
+ )
2698
+
2699
+ if cache_exist == "local":
2700
+ # Allocate cache to cache manager manually to mimick
2701
+ for req in requests:
2702
+ scheduler.encoder_cache_manager.allocate(req, 0)
2703
+ else:
2704
+ # Make sure local encoder cache empty
2705
+ scheduler.encoder_cache_manager.cached = {}
2706
+
2707
+ if cache_exist == "connector_only":
2708
+ # Cache exist in ec_connector
2709
+ scheduler.ec_connector.has_caches = Mock(return_value=[True])
2710
+ elif cache_exist == "no_where":
2711
+ scheduler.ec_connector.has_caches = Mock(return_value=[False])
2712
+
2713
+ output = scheduler.schedule()
2714
+ assert len(output.scheduled_new_reqs) == len(requests)
2715
+ assert output.scheduled_cached_reqs.num_reqs == 0
2716
+ assert len(output.finished_req_ids) == 0
2717
+ for req_id, num_tokens in output.num_scheduled_tokens.items():
2718
+ assert num_tokens == len(requests[int(req_id)].prompt_token_ids)
2719
+
2720
+ ## Encoder-cache-specific checks:
2721
+ # mm_hashes of requests exist in cache after scheduling for all scenario
2722
+ _assert_right_encoder_cache_allocated(scheduler, requests=requests)
2723
+
2724
+ # Should only call update_state_after_alloc when loaded externally
2725
+ if cache_exist == "connector_only":
2726
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(
2727
+ requests[-1], 0
2728
+ )
2729
+
2730
+ # Concat mm_features for the 10 requests together
2731
+ mm_features_list = [feature for req in requests for feature in req.mm_features]
2732
+
2733
+ # Check metadata should contain mm data for all 10 requests
2734
+ _assert_right_ec_connector_metadata(output, mm_features_list=mm_features_list)
2735
+ else:
2736
+ scheduler.ec_connector.update_state_after_alloc.assert_not_called()
2737
+ # ECConnector should carry no metadata
2738
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
2739
+
2740
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2741
+
2742
+ # Should only schedule encoder input when cache is not found anywhere
2743
+ if cache_exist == "no_where":
2744
+ _assert_right_encoder_inputs(
2745
+ output,
2746
+ requests=requests,
2747
+ expected_encoder_inputs=[[0] for _ in range(10)],
2748
+ expected_total_reqs=10,
2749
+ )
2750
+ else:
2751
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2752
+
2753
+
2754
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2755
+def test_ec_connector_unable_to_allocate(use_kv_connector):
2756
+ """
2757
+ Test whether scheduler with ECConnector is able to handle
2758
+ unable to allocate (run out of blocks).
2759
+ """
2760
+
2761
+ # Setup Scheduler With Mock External Cache Hit.
2762
+ BLOCK_SIZE = 4
2763
+ NUM_BLOCKS = 10
2764
+ scheduler = create_scheduler(
2765
+ model="llava-hf/llava-1.5-7b-hf",
2766
+ enable_prefix_caching=True,
2767
+ use_kv_connector=use_kv_connector,
2768
+ block_size=BLOCK_SIZE,
2769
+ num_blocks=NUM_BLOCKS,
2770
+ use_ec_connector=True,
2771
+ ec_role="ec_consumer",
2772
+ )
2773
+
2774
+ # Mock ec_connector load external cache behavior
2775
+ scheduler.ec_connector.has_caches = Mock(return_value=[True])
2776
+ scheduler.ec_connector.update_state_after_alloc = Mock(
2777
+ wraps=scheduler.ec_connector.update_state_after_alloc
2778
+ )
2779
+
2780
+ # Create two requests. The second request will not be able to
2781
+ # allocate slots because it will not have enough blocks.
2782
+ NUM_REQUESTS = 2
2783
+ NUM_TOKENS = (NUM_BLOCKS // 2 + 1) * BLOCK_SIZE
2784
+ MAX_TOKENS = 2
2785
+ requests = create_requests(
2786
+ num_requests=NUM_REQUESTS,
2787
+ num_tokens=NUM_TOKENS,
2788
+ mm_hashes_list=[["hash_1"], ["hash_2"]],
2789
+ mm_positions=[
2790
+ [PlaceholderRange(offset=1, length=10)] for _ in range(NUM_REQUESTS)
2791
+ ],
2792
+ max_tokens=MAX_TOKENS,
2793
+ block_size=BLOCK_SIZE,
2794
+ )
2795
+ req_ids = []
2796
+ req_to_index = {}
2797
+ for i, request in enumerate(requests):
2798
+ scheduler.add_request(request)
2799
+ req_ids.append(request.request_id)
2800
+ req_to_index[request.request_id] = i
2801
+
2802
+ # Setup MODEL_RUNNER_OUTPUT to be run in _step_until_done later
2803
+ MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2804
+ req_ids=req_ids,
2805
+ req_id_to_index=req_to_index,
2806
+ sampled_token_ids=[[1000]] * len(req_ids),
2807
+ logprobs=None,
2808
+ prompt_logprobs_dict={},
2809
+ pooler_output=[],
2810
+ )
2811
+
2812
+ # Just one request should be running.
2813
+ output = scheduler.schedule()
2814
+ scheduled_tokens = output.num_scheduled_tokens[scheduler.running[0].request_id]
2815
+ assert scheduled_tokens == NUM_TOKENS
2816
+ assert len(scheduler.running) == 1
2817
+ assert len(scheduler.waiting) == 1
2818
+
2819
+ # Should have called update_state_after_alloc for external load
2820
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(
2821
+ scheduler.running[0], 0
2822
+ )
2823
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2824
+
2825
+ # All memory should be freed, with one request waiting.
2826
+ _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2827
+ assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
2828
+ assert len(scheduler.running) == 0
2829
+ assert len(scheduler.waiting) == 1
2830
+
2831
+ # Just one request should be running.
2832
+ output = scheduler.schedule()
2833
+ scheduled_tokens = output.num_scheduled_tokens[scheduler.running[0].request_id]
2834
+ assert scheduled_tokens == NUM_TOKENS
2835
+ assert len(scheduler.running) == 1
2836
+ assert len(scheduler.waiting) == 0
2837
+
2838
+ # update_state_after_alloc should be called for loading external cache
2839
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(
2840
+ scheduler.running[0], 0
2841
+ )
2842
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2843
+
2844
+ # All memory should be freed, with no requests waiting / running.
2845
+ _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2846
+ assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
2847
+ assert len(scheduler.running) == 0
2848
+ assert len(scheduler.waiting) == 0
2849
+
2850
+
2851
+@pytest.mark.parametrize("cache_exist", ["local", "connector_only", "no_where"])
2852
+@pytest.mark.parametrize("use_kv_connector", [False, True])
2853
+def test_priority_scheduling_ec_connector_preemption_and_resumption(
2854
+ cache_exist, use_kv_connector
2855
+):
2856
+ """Test that priority scheduling preempts lower priority requests
2857
+ when out of KV cache space."""
2858
+ # Create scheduler with very limited memory to force preemption
2859
+ scheduler = create_scheduler_with_priority(
2860
+ model="llava-hf/llava-1.5-7b-hf",
2861
+ enable_prefix_caching=True,
2862
+ max_num_seqs=2, # allow multiple requests
2863
+ # kv connector should not effect test results
2864
+ use_kv_connector=use_kv_connector,
2865
+ num_blocks=15, # can hold 244 tokens with 14 blocks (first block is null)
2866
+ block_size=16, # standard block size
2867
+ use_ec_connector=True,
2868
+ ec_role="ec_consumer",
2869
+ )
2870
+
2871
+ # Mock cache hit: Both cache exist in connector (at E->PD initially)
2872
+ scheduler.ec_connector.has_caches = Mock(return_value=[True])
2873
+ scheduler.ec_connector.update_state_after_alloc = Mock(
2874
+ wraps=scheduler.ec_connector.update_state_after_alloc
2875
+ )
2876
+
2877
+ # Create a request and schedule it (and to be preempted)
2878
+ request_low = create_requests_with_priority(
2879
+ num_requests=1,
2880
+ priorities=[1],
2881
+ arrival_times=[0.0],
2882
+ num_tokens=94,
2883
+ mm_hashes_list=[["hash_low"]],
2884
+ # NOTE: this test only preempt the last block.
2885
+ # Setting mm_position at the last block can force to recompute encoding
2886
+ mm_positions=[[PlaceholderRange(offset=82, length=10)]],
2887
+ starting_idx=0,
2888
+ )[0]
2889
+ scheduler.add_request(request_low)
2890
+ # 1st schedule
2891
+ output = scheduler.schedule()
2892
+
2893
+ assert len(output.scheduled_new_reqs) == 1
2894
+ scheduled_tokens = output.num_scheduled_tokens[request_low.request_id]
2895
+ assert scheduled_tokens == 94
2896
+ assert len(scheduler.waiting) == 0
2897
+ assert len(scheduler.running) == 1
2898
+
2899
+ ## Encoder-cache-specific checks:
2900
+ # Encoder cache should contain mm items from request
2901
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request_low])
2902
+
2903
+ # Verify update_state_after_alloc called (external load)
2904
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request_low, 0)
2905
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2906
+
2907
+ # ECConnector should carry metadata of request
2908
+ _assert_right_ec_connector_metadata(
2909
+ output, mm_features_list=request_low.mm_features
2910
+ )
2911
+
2912
+ # Scheduled encoder input should be empty; no mm to compute
2913
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2914
+
2915
+ # Simulate model execution - 1st decode
2916
+ model_output = ModelRunnerOutput(
2917
+ req_ids=[request_low.request_id],
2918
+ req_id_to_index={request_low.request_id: 0},
2919
+ sampled_token_ids=[[100]],
2920
+ # spec_token_ids=None,
2921
+ logprobs=None,
2922
+ prompt_logprobs_dict={},
2923
+ pooler_output=[],
2924
+ )
2925
+ scheduler.update_from_output(output, model_output)
2926
+
2927
+ # Create a high priority request and schedule it
2928
+ request_high = create_requests_with_priority(
2929
+ num_requests=1,
2930
+ priorities=[0],
2931
+ arrival_times=[1.0],
2932
+ num_tokens=128,
2933
+ mm_hashes_list=[["hash_high"]],
2934
+ mm_positions=[[PlaceholderRange(offset=1, length=10)]],
2935
+ max_tokens=2,
2936
+ starting_idx=1,
2937
+ )[0]
2938
+ scheduler.add_request(request_high)
2939
+ # 2nd schedule
2940
+ output = scheduler.schedule()
2941
+
2942
+ # KV cache should be full at this point
2943
+ assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 0
2944
+ assert len(output.scheduled_new_reqs) == 1
2945
+ assert output.scheduled_cached_reqs.num_reqs == 1
2946
+ assert len(scheduler.waiting) == 0
2947
+ assert len(scheduler.running) == 2
2948
+
2949
+ ## Encoder-cache-specific checks:
2950
+ # Encoder cache should contain mm items from request
2951
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request_high])
2952
+
2953
+ # Verify update_state_after_alloc called (external load)
2954
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request_high, 0)
2955
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
2956
+
2957
+ # ECConnector should carry metadata of request
2958
+ _assert_right_ec_connector_metadata(
2959
+ output, mm_features_list=request_high.mm_features
2960
+ )
2961
+
2962
+ # Scheduled encoder input should be empty; no mm to compute
2963
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2964
+
2965
+ # Simulate model execution - 2nd decode
2966
+ requests = [request_low, request_high]
2967
+ model_output = ModelRunnerOutput(
2968
+ req_ids=[req.request_id for req in requests],
2969
+ req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
2970
+ sampled_token_ids=[[100] for _ in requests],
2971
+ # spec_token_ids=None,
2972
+ logprobs=None,
2973
+ prompt_logprobs_dict={},
2974
+ pooler_output=[],
2975
+ )
2976
+ scheduler.update_from_output(output, model_output)
2977
+
2978
+ # 3rd schedule - - this should trigger preemption
2979
+ # req_low needs 96 tokens = 6 blocks
2980
+ # req_high needs 129 tokens = 9 blocks
2981
+ # so doesn't fit in 14 blocks.
2982
+ output = scheduler.schedule()
2983
+
2984
+ # Should have preempted req_low
2985
+ assert len(output.scheduled_new_reqs) == 0
2986
+ assert output.scheduled_cached_reqs.num_reqs == 1
2987
+ assert output.scheduled_cached_reqs.req_ids[0] == request_high.request_id
2988
+ assert scheduler.requests[request_low.request_id].status == RequestStatus.PREEMPTED
2989
+ assert len(scheduler.waiting) == 1
2990
+ assert len(scheduler.running) == 1
2991
+
2992
+ ## Encoder-cache-specific checks:
2993
+ # request_high is in decode phase now
2994
+ # ECConnector should carry no metadata
2995
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
2996
+
2997
+ # Scheduled encoder input should be empty; no mm to compute
2998
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
2999
+
3000
+ # Simulate model execution - 3rd decode, after req_low was preempted
3001
+ requests = [request_low, request_high]
3002
+ model_output = ModelRunnerOutput(
3003
+ req_ids=[req.request_id for req in requests],
3004
+ req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
3005
+ sampled_token_ids=[[100], [100, 200]],
3006
+ # spec_token_ids=None,
3007
+ logprobs=None,
3008
+ prompt_logprobs_dict={},
3009
+ pooler_output=[],
3010
+ )
3011
+ # Finish the requests to make room for the preempted requests to resume
3012
+ # req_high is finished after outputing 2 tokens
3013
+ scheduler.update_from_output(output, model_output)
3014
+ scheduler.finish_requests(
3015
+ request_high.request_id, RequestStatus.FINISHED_LENGTH_CAPPED
3016
+ )
3017
+
3018
+ # Set up to test different encoder cache exsistence scenario after preemption
3019
+ # Order of getting encoder cache should be: local cache -> connector-> compute
3020
+ # By default, the cache should still exist in local in this test case
3021
+ if cache_exist != "local":
3022
+ # Make local encoder cache empty
3023
+ scheduler.encoder_cache_manager.cached = {}
3024
+
3025
+ if cache_exist == "connector_only":
3026
+ # Cache exist in ec_connector
3027
+ scheduler.ec_connector.has_caches = Mock(return_value=[True])
3028
+ elif cache_exist == "no_where":
3029
+ scheduler.ec_connector.has_caches = Mock(return_value=[False])
3030
+
3031
+ # 4th Schedule - this should trigger req_low resumption from waiting
3032
+ output = scheduler.schedule()
3033
+ scheduled_cached_reqs = output.scheduled_cached_reqs
3034
+ resumed_from_preemption = scheduled_cached_reqs.resumed_from_preemption
3035
+
3036
+ assert len(output.scheduled_new_reqs) == 0
3037
+ assert scheduled_cached_reqs.num_reqs == 1
3038
+ assert len(scheduler.waiting) == 0
3039
+ assert len(scheduler.running) == 1
3040
+
3041
+ # Preempted request resumed in scheduled_cached_reqs
3042
+ assert len(resumed_from_preemption) == 1
3043
+ assert len(scheduled_cached_reqs.resumed_req_token_ids) == 1
3044
+ assert resumed_from_preemption[0]
3045
+ assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
3046
+ assert scheduled_cached_reqs.resumed_req_token_ids[0] is not None
3047
+ ## Resumed tokens include 94 prompt tokens and 2 decoded tokens
3048
+ assert len(scheduled_cached_reqs.resumed_req_token_ids[0]) == 96
3049
+ assert scheduled_cached_reqs.resumed_req_token_ids[0][95] == 100
3050
+ assert scheduler.running[0].request_id == request_low.request_id
3051
+ assert request_high.request_id in output.finished_req_ids
3052
+
3053
+ ## Encoder-cache-specific checks:
3054
+ # mm_hash of request_low exists in cache after scheduling for all scenario
3055
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request_low])
3056
+
3057
+ # Should only call update_state_after_alloc when loaded externally
3058
+ if cache_exist == "connector_only":
3059
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(
3060
+ request_low, 0
3061
+ )
3062
+ _assert_right_ec_connector_metadata(
3063
+ output, mm_features_list=request_low.mm_features
3064
+ )
3065
+ else:
3066
+ scheduler.ec_connector.update_state_after_alloc.assert_not_called()
3067
+ # ECConnector should carry no metadata
3068
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
3069
+
3070
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
3071
+
3072
+ # Should only schedule encoder input when cache is not found anywhere
3073
+ if cache_exist == "no_where":
3074
+ _assert_right_encoder_inputs(
3075
+ output,
3076
+ requests=[request_low],
3077
+ expected_encoder_inputs=[[0]],
3078
+ expected_total_reqs=1,
3079
+ )
3080
+ else:
3081
+ _assert_right_encoder_inputs(output, expected_total_reqs=0)
3082
+
3083
+
3084
+@pytest.mark.parametrize("use_kv_connector", [False, True])
3085
+def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connector):
3086
+ """
3087
+ Scenario:
3088
+ - Encoder cache size: 32
3089
+ - Request A: 1 feature (12 tokens) → NOT cached remotely.
3090
+ - Request B: 3 features (3 x 10 tokens) → ALL cached remotely.
3091
+
3092
+ Steps:
3093
+ 1. Schedule Request A (locally uses 12 tokens).
3094
+ 2. Schedule Request B (remote cache) - only schedule 1st and 2nd
3095
+ 3. Free A's cache, then schedule B again (continuation) - schedule 3rd image
3096
+ """
3097
+ scheduler = create_scheduler(
3098
+ model="llava-hf/llava-1.5-7b-hf",
3099
+ max_num_batched_tokens=1024,
3100
+ enable_prefix_caching=True,
3101
+ use_kv_connector=use_kv_connector,
3102
+ block_size=16,
3103
+ num_blocks=11, # Can hold 160 tokens (first block is null)
3104
+ use_ec_connector=True,
3105
+ ec_role="ec_consumer",
3106
+ )
3107
+
3108
+ # Limit the number of availiable slots of EncoderCacheManager
3109
+ scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=32)
3110
+
3111
+ # Create MM request1
3112
+ NUM_TOKENS_1 = 50 # NOTE: includes mm tokens
3113
+ NUM_ENCODER_TOKENS_1 = 12
3114
+ mm_hashes_list_1 = [["hash1_1"]]
3115
+ mm_positions_1 = [[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_1)]]
3116
+
3117
+ request1 = create_requests(
3118
+ num_requests=1,
3119
+ num_tokens=NUM_TOKENS_1,
3120
+ mm_hashes_list=mm_hashes_list_1,
3121
+ mm_positions=mm_positions_1,
3122
+ max_tokens=1, # For simplicity
3123
+ req_ids=["req1"],
3124
+ )[0]
3125
+
3126
+ # Create MM request1 with 3 MM items
3127
+ NUM_TOKENS_2 = 40
3128
+ NUM_ENCODER_TOKENS_2 = 10
3129
+ mm_hashes_list_2 = [["hash2_1", "hash2_2", "hash2_3"]]
3130
+ mm_positions_2 = [
3131
+ [
3132
+ PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_2),
3133
+ PlaceholderRange(offset=12, length=NUM_ENCODER_TOKENS_2),
3134
+ PlaceholderRange(offset=24, length=NUM_ENCODER_TOKENS_2),
3135
+ ]
3136
+ ]
3137
+
3138
+ request2 = create_requests(
3139
+ num_requests=1,
3140
+ num_tokens=NUM_TOKENS_2,
3141
+ mm_hashes_list=mm_hashes_list_2,
3142
+ mm_positions=mm_positions_2,
3143
+ max_tokens=10,
3144
+ req_ids=["req2"],
3145
+ )[0]
3146
+
3147
+ # Mock cache hit: MM of request1 NOT cached remotely, request2 cached remotely
3148
+ scheduler.ec_connector.has_caches = Mock(
3149
+ side_effect=lambda req: [True, True, True] if req == request2 else [False]
3150
+ )
3151
+ scheduler.ec_connector.update_state_after_alloc = Mock(
3152
+ wraps=scheduler.ec_connector.update_state_after_alloc
3153
+ )
3154
+
3155
+ scheduler.add_request(request1)
3156
+ scheduler.add_request(request2)
3157
+ output = scheduler.schedule()
3158
+
3159
+ # Now, since encoder cache manager can only store 32 tokens
3160
+ # It should allocated mm item hash1_1, hash2_1 and hash2_2
3161
+ scheduled_tokens = output.num_scheduled_tokens[request1.request_id]
3162
+ assert scheduled_tokens == NUM_TOKENS_1
3163
+ assert scheduler.get_num_unfinished_requests() == 2
3164
+
3165
+ # Encoder cache should contain mm item from request1
3166
+ _assert_right_encoder_cache_allocated(
3167
+ scheduler, hashes_to_check=["hash1_1", "hash2_1", "hash2_2"]
3168
+ )
3169
+
3170
+ # request2's 2nd mm item is the last call of update_state_after_alloc
3171
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request2, 1)
3172
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
3173
+
3174
+ # ECConnector should carry metadata of hash2_1 and hash2_2 ONLY
3175
+ _assert_right_ec_connector_metadata(
3176
+ output, mm_features_list=[request2.mm_features[0], request2.mm_features[1]]
3177
+ )
3178
+
3179
+ # Should schedule ONLY 1 encoder input
3180
+ _assert_right_encoder_inputs(
3181
+ output,
3182
+ requests=[request1],
3183
+ expected_encoder_inputs=[[0]], # index 0 of the mm item of request1
3184
+ expected_total_reqs=1,
3185
+ )
3186
+
3187
+ # Simulate model execution 1 step
3188
+ model_output = ModelRunnerOutput(
3189
+ req_ids=[request1.request_id, request2.request_id],
3190
+ req_id_to_index={request1.request_id: 0, request2.request_id: 1},
3191
+ sampled_token_ids=[[100], [121]],
3192
+ # spec_token_ids=None,
3193
+ logprobs=None,
3194
+ prompt_logprobs_dict={},
3195
+ pooler_output=[],
3196
+ )
3197
+ scheduler.update_from_output(output, model_output)
3198
+
3199
+ # request1 is finished after outputing 1 token
3200
+ # Finish request
3201
+ scheduler.finish_requests(request1.request_id, RequestStatus.FINISHED_LENGTH_CAPPED)
3202
+ assert scheduler.get_num_unfinished_requests() == 1
3203
+
3204
+ # Schedule again; Now request1's encoder cache should be freed
3205
+ # -> hash2_3 can be scheduled and allocated
3206
+ output = scheduler.schedule()
3207
+
3208
+ # Check
3209
+ # Should schedule all tokens
3210
+ scheduled_tokens = output.num_scheduled_tokens[request2.request_id]
3211
+ print(f"Hero: scheduled_tokens for req2: {scheduled_tokens}")
3212
+ print(f"hero: num_scheduled_tokens 2: {output.num_scheduled_tokens}")
3213
+
3214
+ # Encoder cache should contain all mm items from request2
3215
+ _assert_right_encoder_cache_allocated(scheduler, requests=[request2])
3216
+
3217
+ # request2's 3rd mm item is the ONLY call of update_state_after_alloc
3218
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(request2, 2)
3219
+ scheduler.ec_connector.update_state_after_alloc.assert_called_once()
3220
+
3221
+ scheduler.ec_connector.update_state_after_alloc.reset_mock()
3222
+
3223
+ # ECConnector should carry metadata for hash2_3 ONLY
3224
+ _assert_right_ec_connector_metadata(
3225
+ output, mm_features_list=[request2.mm_features[2]]
3226
+ )
3227
+
3228
+ # Should schedule no encoder input
3229
+ _assert_right_encoder_inputs(
3230
+ output,
3231
+ expected_total_reqs=0,
3232
+ )
3233
+
3234
+
3235
+# ==============================================================================
3236
+# EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests end
3237
+# ==============================================================================
tests/v1/core/utils.py
+62
-13
@@ -5,6 +5,7 @@ import torch
5
6
from vllm.config import (
7
CacheConfig,
8
+ ECTransferConfig,
9
KVTransferConfig,
10
ModelConfig,
11
SchedulerConfig,
@@ -46,6 +47,8 @@ def create_scheduler(
47
num_speculative_tokens: int | None = None,
48
skip_tokenizer_init: bool = False,
49
async_scheduling: bool = False,
50
+ use_ec_connector: bool = False,
51
+ ec_role: str | None = None,
52
) -> Scheduler | AsyncScheduler:
53
"""Create scheduler under test.
54
@@ -107,12 +110,23 @@ def create_scheduler(
110
model="ngram", num_speculative_tokens=num_speculative_tokens
111
)
112
113
+ ec_transfer_config = (
114
+ ECTransferConfig(
115
+ ec_connector="ECSharedStorageConnector",
116
+ ec_role=ec_role,
117
+ ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test"},
118
+ )
119
+ if use_ec_connector
120
+ else None
121
+ )
122
+
123
vllm_config = VllmConfig(
124
scheduler_config=scheduler_config,
125
model_config=model_config,
126
cache_config=cache_config,
127
kv_transfer_config=kv_transfer_config,
128
speculative_config=speculative_config,
129
+ ec_transfer_config=ec_transfer_config,
130
)
131
kv_cache_config = KVCacheConfig(
132
num_blocks=num_blocks, # A large number of blocks to hold all requests
@@ -140,12 +154,14 @@ _none_hash_initialized = False
154
def create_requests(
155
num_requests: int,
156
num_tokens: int = 10,
157
+ mm_hashes_list: list[list[str]] | None = None,
158
mm_positions: list[list[PlaceholderRange]] | None = None,
159
max_tokens: int = 16,
160
stop_token_ids: list[int] | None = None,
161
prompt_logprobs: int | None = None,
162
same_prompt: bool = False,
163
block_size: int = 16,
164
+ req_ids: list[str] | None = None,
165
) -> list[Request]:
166
global _none_hash_initialized
167
if not _none_hash_initialized:
@@ -160,25 +176,58 @@ def create_requests(
176
prompt_logprobs=prompt_logprobs,
177
)
178
requests = []
179
+
180
+ if mm_hashes_list is not None:
181
+ # NOTE: allow manual input; some mm items can have the same identifier
182
+ # no. of mm_hashes and mm_positions for each request should be identical
183
+ assert mm_positions is not None, (
184
+ "mm_positions must be provided when mm_hashes_list is provided"
185
+ )
186
+ assert len(mm_hashes_list) == len(mm_positions) == num_requests
187
+ assert [len(h) for h in mm_hashes_list] == [len(p) for p in mm_positions]
188
+
189
+ # Since same identifier would imply they are identical encoder output
190
+ # Verify mm items with identical identifier are having mm_position.length
191
+ seen_hashes: dict[str, int] = {}
192
+
193
+ if req_ids:
194
+ assert len(req_ids) == num_requests
195
+ else:
196
+ req_ids = [f"{i}" for i in range(num_requests)]
197
+
198
for i in range(num_requests):
199
mm_features = []
165
- if mm_positions is not None:
166
- mm_position = mm_positions[i]
167
- for j, position in enumerate(mm_position):
168
- # Dummy hash for each mm item should be unique
169
- # since encoder cache tracks entries by hash
200
+
201
+ for j, position in enumerate(
202
+ mm_positions[i] if mm_positions is not None else []
203
+ ):
204
+ if mm_hashes_list is not None:
205
+ identifier = mm_hashes_list[i][j]
206
+
207
+ # Verify if position length is identical
208
+ position_length = position.length
209
+ if identifier in seen_hashes:
210
+ assert seen_hashes[identifier] == position_length, (
211
+ f"mm_hash '{identifier}' has inconsistent position lengths: "
212
+ f"previously {seen_hashes[identifier]}, now {position_length} "
213
+ f"at request {i}, position {j}"
214
+ )
215
+ else:
216
+ seen_hashes[identifier] = position_length
217
+ else:
218
+ # Unique dummy hash for each mm item
219
identifier = f"hash{i}_{j}"
171
- mm_feature = MultiModalFeatureSpec(
172
- data=MultiModalKwargsItem.dummy("dummy_m"),
173
- mm_position=position,
174
- identifier=identifier,
175
- modality="image",
176
- )
177
- mm_features.append(mm_feature)
220
+ mm_feature = MultiModalFeatureSpec(
221
+ data=MultiModalKwargsItem.dummy("dummy_m"),
222
+ mm_position=position,
223
+ identifier=identifier,
224
+ modality="image",
225
+ )
226
+ mm_features.append(mm_feature)
227
228
prompt_token_ids = [0] * num_tokens if same_prompt else [i] * num_tokens
229
request = Request(
181
- request_id=f"{i}",
230
+ request_id=req_ids[i],
231
prompt_token_ids=prompt_token_ids,
232
sampling_params=sampling_params,
233
pooling_params=None,
tests/v1/ec_connector/integration/README.md
new
+171
@@ -0,0 +1,171 @@
1
+# EPD Correctness Test
2
+
3
+This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance.
4
+
5
+## What It Tests
6
+
7
+- **Baseline**: Single vLLM instance serving a multimodal model
8
+- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance
9
+- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance
10
+- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance
11
+
12
+The test ensures that disaggregated encoding produces **identical** outputs to the baseline.
13
+
14
+Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D
15
+
16
+Please refer to [Disaggregated Encoder Feature](../../../docs/features/disagg_encoder.md) for the detailed explanation for the EPD features.
17
+
18
+## Files
19
+
20
+- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests)
21
+- `test_epd_correctness.py` - Python test script (compares outputs)
22
+
23
+## Usage
24
+
25
+### Multimodal Prompts (Default)
26
+
27
+```bash
28
+cd vllm
29
+./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
30
+```
31
+
32
+This runs the test with actual multimodal (image) prompts.
33
+
34
+### Text-Only Prompts
35
+
36
+```bash
37
+cd vllm
38
+USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
39
+```
40
+
41
+This runs a quick test with text-only prompts to verify the setup works.
42
+
43
+### Custom Configuration
44
+
45
+```bash
46
+# Use specific GPUs
47
+GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
48
+
49
+# Use specific ports
50
+ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
51
+
52
+# Use specific model
53
+MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
54
+
55
+# Use specific storage path
56
+EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
57
+```
58
+
59
+## How It Works
60
+
61
+### Step 1: Baseline
62
+
63
+1. Start single vLLM instance on GPU
64
+2. Run test prompts (multimodal or text-only)
65
+3. Save outputs to `.vllm_epd_baseline.txt`
66
+4. Shutdown instance
67
+
68
+### Step 2: EPD (1E + 1PD)
69
+
70
+1. Clear encoder cache storage
71
+2. Start instances and proxy
72
+3. Run same test prompts
73
+4. Assert outputs match baseline exactly
74
+5. Shutdown instances
75
+
76
+### Step 3: EPD (1E + 1P + 1D)
77
+
78
+1. Clear encoder cache storage
79
+2. Start instances and proxy
80
+3. Run same test prompts
81
+4. Assert outputs match baseline exactly
82
+5. Shutdown instances
83
+
84
+## Test Scenarios
85
+
86
+### Multimodal Prompts (--use_mm_prompts)
87
+
88
+Tests encoder cache transfer:
89
+
90
+- Single image query
91
+- Multiple images in one request
92
+- Mixed image and text
93
+- Image with detailed questions
94
+
95
+### Text-Only Prompts (default)
96
+
97
+Quick sanity check:
98
+
99
+- Simple text queries
100
+- Text-only explanations
101
+- Verifies proxy routing works
102
+
103
+## Expected Behavior
104
+
105
+### ✅ Test Passes When
106
+
107
+- All disagg outputs match baseline outputs exactly
108
+- No errors during instance startup
109
+- Encoder cache is properly saved and loaded
110
+- Proxy correctly routes requests
111
+
112
+### ❌ Test Fails When
113
+
114
+- Outputs differ between baseline and disagg
115
+- Server startup fails
116
+- Encoder cache not found (should fallback to local execution)
117
+- Proxy routing errors
118
+
119
+## Notes
120
+
121
+- The test uses deterministic generation (`temperature=0.0`, `seed=42`)
122
+- Encoder cache should enable exact output reproduction
123
+- Test cleans up all instances and cache files after completion
124
+- Safe to run multiple times (idempotent)
125
+- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/online_serving/disaggregated_encoder/README.md`
126
+
127
+## Requirements
128
+
129
+- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline)
130
+ - 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now.
131
+- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct)
132
+- Internet access (for accessing vllm test images)
133
+
134
+## Debugging
135
+
136
+### Check Logs
137
+
138
+Logs and baseline output are saved in `/tmp/` by default.
139
+Can be customized by changing the environment variables.
140
+
141
+### Check Encoder Cache
142
+
143
+```bash
144
+# Verify cache files are created
145
+ls -la $EC_SHARED_STORAGE_PATH/
146
+
147
+# Should see directories with mm_hash names
148
+# Each containing encoder_cache.safetensors
149
+```
150
+
151
+### Manual Testing
152
+
153
+Run individual components:
154
+
155
+```bash
156
+# Baseline only
157
+python test_epd_correctness.py \
158
+ --service_url http://localhost:8000 \
159
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
160
+ --mode baseline \
161
+ --baseline_file test_output.txt \
162
+ --use_mm_prompts
163
+
164
+# Disagg only (requires baseline output file!)
165
+python test_epd_correctness.py \
166
+ --service_url http://localhost:8000 \
167
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
168
+ --mode disagg \
169
+ --baseline_file test_output.txt \
170
+ --use_mm_prompts
171
+```
tests/v1/ec_connector/integration/hato.jpg
Binary files /dev/null and b/tests/v1/ec_connector/integration/hato.jpg differ
tests/v1/ec_connector/integration/run_epd_correctness_test.sh
new
+476
@@ -0,0 +1,476 @@
1
+#!/bin/bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# EPD (Encoder-Prefill-Decode) Correctness Test
6
+#
7
+# This script tests that EPD disaggregation produces the same outputs as baseline.
8
+# It runs:
9
+# 1. Baseline: Single vLLM instance
10
+# 2. EPD: 1E + 1PD setup
11
+# 3. Baseline for (E + P + D): 1P + 1D vLLM instances disagg
12
+# 4. EPD: 1E + 1P + 1D setup
13
+
14
+# For GPU usage
15
+
16
+# set -xe
17
+
18
+# Find the git repository root directory
19
+GIT_ROOT=$(git rev-parse --show-toplevel)
20
+
21
+# Model to test
22
+MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
23
+
24
+# Set 1 to use multimodal prompts; else to use text-only
25
+USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
26
+MM_FLAG=""
27
+if [ $USE_MM_PROMPTS = "1" ]; then
28
+ MM_FLAG="--use_mm_prompts"
29
+fi
30
+
31
+# GPU configuration
32
+GPU_E="${GPU_E:-0}"
33
+GPU_P="${GPU_P:-1}"
34
+GPU_D="${GPU_D:-2}"
35
+GPU_SINGLE="${GPU_SINGLE:-$GPU_P}"
36
+GPU_PD="${GPU_PD:-$GPU_P}"
37
+
38
+# Port
39
+ENCODE_PORT="${ENCODE_PORT:-19534}"
40
+PREFILL_PORT="${PREFILL_PORT:-19535}"
41
+DECODE_PORT="${DECODE_PORT:-19536}"
42
+PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}"
43
+ENDPOINT_PORT="${ENDPOINT_PORT:-10001}"
44
+
45
+# Storage path for encoder cache
46
+EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache_test}"
47
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-600}"
48
+
49
+# Output file for baseline comparison and logs
50
+LOG_PATH="${LOG_PATH:-/tmp}"
51
+BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_baseline.txt}"
52
+BASELINE_PD_FILE="${BASELINE_PD_FILE:-/tmp/vllm_epd_baseline.txt}"
53
+
54
+mkdir -p $LOG_PATH
55
+
56
+# Trap the SIGINT signal (triggered by Ctrl+C)
57
+trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
58
+
59
+# Wait for server to be ready
60
+wait_for_server() {
61
+ local port=$1
62
+ timeout "$TIMEOUT_SECONDS" bash -c "
63
+ until curl -s localhost:${port}/v1/chat/completions > /dev/null; do
64
+ sleep 1
65
+ done" && return 0 || return 1
66
+}
67
+
68
+# Cleanup function
69
+cleanup_instances() {
70
+ echo "Cleaning up any running vLLM instances..."
71
+ pkill -f "vllm serve" || true
72
+ pkill -f "disagg_epd_proxy.py" || true
73
+ sleep 2
74
+}
75
+
76
+# Function to run baseline (single instance)
77
+run_baseline() {
78
+ echo "================================"
79
+ echo "Running BASELINE (single instance)"
80
+ echo "================================"
81
+
82
+ cleanup_instances
83
+ rm -rf "$EC_SHARED_STORAGE_PATH"
84
+
85
+ local PORT=$ENDPOINT_PORT
86
+
87
+ # Start baseline instance
88
+ echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT"
89
+ CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \
90
+ --port $PORT \
91
+ --enforce-eager \
92
+ --gpu-memory-utilization 0.7 \
93
+ --max-num-seqs 128 \
94
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
95
+ > $LOG_PATH/baseline.log 2>&1 &
96
+
97
+ local BASELINE_PID=$!
98
+
99
+ # Wait for baseline to start
100
+ echo "Waiting for baseline instance to start..."
101
+ wait_for_server $PORT
102
+
103
+ curl http://127.0.0.1:$PORT/v1/models
104
+ echo ""
105
+
106
+ # Run test in baseline mode
107
+ echo "Running baseline..."
108
+
109
+ python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
110
+ --service_url "http://localhost:$PORT" \
111
+ --model_name "$MODEL" \
112
+ --mode baseline \
113
+ --baseline_file "$BASELINE_FILE" \
114
+ $MM_FLAG
115
+
116
+ # Cleanup baseline
117
+ echo "Stopping baseline instance..."
118
+ kill $BASELINE_PID 2>/dev/null || true
119
+ sleep 2
120
+ cleanup_instances
121
+}
122
+
123
+# Function to run EPD with 1E + 1PD
124
+run_epd_1e_1pd() {
125
+ echo "================================"
126
+ echo "Running EPD (1E + 1PD)"
127
+ echo "================================"
128
+
129
+ cleanup_instances
130
+ rm -rf "$EC_SHARED_STORAGE_PATH"
131
+ mkdir -p "$EC_SHARED_STORAGE_PATH"
132
+
133
+ local ENCODE_PORT=$ENCODE_PORT
134
+ local PREFILL_DECODE_PORT=$PREFILL_DECODE_PORT
135
+ local PROXY_PORT=$ENDPOINT_PORT
136
+
137
+ declare -a PIDS=()
138
+
139
+ # Start encoder instance
140
+ echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
141
+ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
142
+ --port $ENCODE_PORT \
143
+ --enforce-eager \
144
+ --gpu-memory-utilization 0.01 \
145
+ --enable-request-id-headers \
146
+ --no-enable-prefix-caching \
147
+ --max-num-batched-tokens 114688 \
148
+ --max-num-seqs 128 \
149
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
150
+ --ec-transfer-config '{
151
+ "ec_connector": "ECSharedStorageConnector",
152
+ "ec_role": "ec_producer",
153
+ "ec_connector_extra_config": {
154
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
155
+ }
156
+ }' \
157
+ > $LOG_PATH/1e1pd_encoder.log 2>&1 &
158
+ PIDS+=($!)
159
+
160
+ # Start prefill+decode instance
161
+ echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT"
162
+ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
163
+ --port $PREFILL_DECODE_PORT \
164
+ --enforce-eager \
165
+ --gpu-memory-utilization 0.7 \
166
+ --enable-request-id-headers \
167
+ --max-num-seqs 128 \
168
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
169
+ --ec-transfer-config '{
170
+ "ec_connector": "ECSharedStorageConnector",
171
+ "ec_role": "ec_consumer",
172
+ "ec_connector_extra_config": {
173
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
174
+ }
175
+ }' \
176
+ > $LOG_PATH/1e1pd_pd.log 2>&1 &
177
+ PIDS+=($!)
178
+
179
+ # Wait for instances to start
180
+ echo "Waiting for encoder instance..."
181
+ wait_for_server $ENCODE_PORT
182
+ echo "Waiting for PD instance..."
183
+ wait_for_server $PREFILL_DECODE_PORT
184
+
185
+ # Start proxy
186
+ echo "Starting EPD proxy on port $PROXY_PORT"
187
+ python "${GIT_ROOT}/examples/online_serving/disaggregated_encoder/disagg_epd_proxy.py" \
188
+ --host "0.0.0.0" \
189
+ --port $PROXY_PORT \
190
+ --encode-servers-urls "http://localhost:$ENCODE_PORT" \
191
+ --prefill-servers-urls "disable" \
192
+ --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
193
+ > $LOG_PATH/1e1pd_proxy.log 2>&1 &
194
+ PIDS+=($!)
195
+
196
+ # Wait for proxy
197
+ echo "Waiting for proxy..."
198
+ wait_for_server $PROXY_PORT
199
+
200
+ curl http://127.0.0.1:$PROXY_PORT/v1/models
201
+ curl http://127.0.0.1:$PROXY_PORT/health
202
+ echo ""
203
+
204
+ echo "All EPD (1E+1PD) services are up!"
205
+
206
+ # Run test in disagg mode
207
+ echo "Running EPD (1E+1PD) correctness test..."
208
+
209
+ python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
210
+ --service_url "http://localhost:$PROXY_PORT" \
211
+ --model_name "$MODEL" \
212
+ --mode disagg \
213
+ --baseline_file "$BASELINE_FILE" \
214
+ $MM_FLAG
215
+
216
+ # Cleanup
217
+ echo "✓✓ 1E+1PD Correctness Test finished"
218
+ echo "Stopping EPD (1E+1PD) instances..."
219
+ for pid in "${PIDS[@]}"; do
220
+ kill $pid 2>/dev/null || true
221
+ done
222
+ sleep 2
223
+ cleanup_instances
224
+}
225
+
226
+# Function to run baseline for 1E + 1P + 1D (PD disagg)
227
+run_baseline_1p_1d() {
228
+ echo "================================"
229
+ echo "Running PD BASELINE (1P + 1D)"
230
+ echo "================================"
231
+
232
+ cleanup_instances
233
+ rm -rf "$EC_SHARED_STORAGE_PATH"
234
+ mkdir -p "$EC_SHARED_STORAGE_PATH"
235
+
236
+ local PREFILL_PORT=$PREFILL_PORT
237
+ local DECODE_PORT=$DECODE_PORT
238
+ local PROXY_PORT=$ENDPOINT_PORT
239
+
240
+ declare -a PIDS=()
241
+
242
+ # Start prefill instance
243
+ echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
244
+ CUDA_VISIBLE_DEVICES="$GPU_P" \
245
+ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
246
+ vllm serve "$MODEL" \
247
+ --port $PREFILL_PORT \
248
+ --enforce-eager \
249
+ --gpu-memory-utilization 0.7 \
250
+ --enable-request-id-headers \
251
+ --max-num-seqs 128 \
252
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
253
+ --kv-transfer-config '{
254
+ "kv_connector": "NixlConnector",
255
+ "kv_role": "kv_producer"
256
+ }' \
257
+ > $LOG_PATH/1p1d_prefill.log 2>&1 &
258
+ PIDS+=($!)
259
+
260
+ # Start decode instance
261
+ echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
262
+ CUDA_VISIBLE_DEVICES="$GPU_D" \
263
+ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
264
+ vllm serve "$MODEL" \
265
+ --port $DECODE_PORT \
266
+ --enforce-eager \
267
+ --gpu-memory-utilization 0.7 \
268
+ --enable-request-id-headers \
269
+ --max-num-seqs 128 \
270
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
271
+ --kv-transfer-config '{
272
+ "kv_connector": "NixlConnector",
273
+ "kv_role": "kv_consumer"
274
+ }' \
275
+ > $LOG_PATH/1p1d_decode.log 2>&1 &
276
+ PIDS+=($!)
277
+
278
+ # Wait for instances to start
279
+ echo "Waiting for prefill instance..."
280
+ wait_for_server $PREFILL_PORT
281
+ echo "Waiting for decode instance..."
282
+ wait_for_server $DECODE_PORT
283
+
284
+ # Start proxy
285
+ echo "Starting EPD proxy on port $PROXY_PORT"
286
+ python "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
287
+ --host "0.0.0.0" \
288
+ --port $PROXY_PORT \
289
+ --prefiller-ports $PREFILL_PORT \
290
+ --decoder-ports $DECODE_PORT \
291
+ > $LOG_PATH/1p1d_proxy.log 2>&1 &
292
+ PIDS+=($!)
293
+
294
+ # Wait for proxy
295
+ echo "Waiting for proxy..."
296
+ wait_for_server $PROXY_PORT
297
+
298
+ curl http://127.0.0.1:$PROXY_PORT/healthcheck
299
+ echo ""
300
+
301
+ echo "All PD (1P+1D) services are up!"
302
+
303
+ # Run test in baseline mode
304
+ echo "Running PD disagg baseline..."
305
+
306
+ python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
307
+ --service_url "http://localhost:$PROXY_PORT" \
308
+ --model_name "$MODEL" \
309
+ --mode baseline_pd \
310
+ --baseline_file "$BASELINE_PD_FILE" \
311
+ $MM_FLAG
312
+
313
+ # Cleanup
314
+ echo "Stopping PD (1P+1D) instances..."
315
+ for pid in "${PIDS[@]}"; do
316
+ kill $pid 2>/dev/null || true
317
+ done
318
+ sleep 2
319
+ cleanup_instances
320
+}
321
+
322
+# Function to run EPD with 1E + 1P + 1D
323
+run_epd_1e_1p_1d() {
324
+ echo "================================"
325
+ echo "Running EPD (1E + 1P + 1D)"
326
+ echo "================================"
327
+
328
+ cleanup_instances
329
+ rm -rf "$EC_SHARED_STORAGE_PATH"
330
+ mkdir -p "$EC_SHARED_STORAGE_PATH"
331
+
332
+ local ENCODE_PORT=$ENCODE_PORT
333
+ local PREFILL_PORT=$PREFILL_PORT
334
+ local DECODE_PORT=$DECODE_PORT
335
+ local PROXY_PORT=$ENDPOINT_PORT
336
+
337
+ declare -a PIDS=()
338
+
339
+ # Start encoder instance
340
+ echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
341
+ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
342
+ --port $ENCODE_PORT \
343
+ --enforce-eager \
344
+ --gpu-memory-utilization 0.01 \
345
+ --enable-request-id-headers \
346
+ --no-enable-prefix-caching \
347
+ --max-num-batched-tokens 114688 \
348
+ --max-num-seqs 128 \
349
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
350
+ --ec-transfer-config '{
351
+ "ec_connector": "ECSharedStorageConnector",
352
+ "ec_role": "ec_producer",
353
+ "ec_connector_extra_config": {
354
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
355
+ }
356
+ }' \
357
+ > $LOG_PATH/1e1p1d_encoder.log 2>&1 &
358
+ PIDS+=($!)
359
+
360
+ # Start prefill instance
361
+ echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
362
+ CUDA_VISIBLE_DEVICES="$GPU_P" \
363
+ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
364
+ vllm serve "$MODEL" \
365
+ --port $PREFILL_PORT \
366
+ --enforce-eager \
367
+ --gpu-memory-utilization 0.7 \
368
+ --enable-request-id-headers \
369
+ --max-num-seqs 128 \
370
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
371
+ --ec-transfer-config '{
372
+ "ec_connector": "ECSharedStorageConnector",
373
+ "ec_role": "ec_consumer",
374
+ "ec_connector_extra_config": {
375
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
376
+ }
377
+ }' \
378
+ --kv-transfer-config '{
379
+ "kv_connector": "NixlConnector",
380
+ "kv_role": "kv_producer"
381
+ }' \
382
+ > $LOG_PATH/1e1p1d_prefill.log 2>&1 &
383
+ PIDS+=($!)
384
+
385
+ # Start decode instance
386
+ echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
387
+ CUDA_VISIBLE_DEVICES="$GPU_D" \
388
+ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
389
+ vllm serve "$MODEL" \
390
+ --port $DECODE_PORT \
391
+ --enforce-eager \
392
+ --gpu-memory-utilization 0.7 \
393
+ --enable-request-id-headers \
394
+ --max-num-seqs 128 \
395
+ --allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
396
+ --kv-transfer-config '{
397
+ "kv_connector": "NixlConnector",
398
+ "kv_role": "kv_consumer"
399
+ }' \
400
+ > $LOG_PATH/1e1p1d_decode.log 2>&1 &
401
+ PIDS+=($!)
402
+
403
+ # Wait for instances to start
404
+ echo "Waiting for encoder instance..."
405
+ wait_for_server $ENCODE_PORT
406
+ echo "Waiting for prefill instance..."
407
+ wait_for_server $PREFILL_PORT
408
+ echo "Waiting for decode instance..."
409
+ wait_for_server $DECODE_PORT
410
+
411
+ # Start proxy
412
+ echo "Starting EPD proxy on port $PROXY_PORT"
413
+ python "${GIT_ROOT}/examples/online_serving/disaggregated_encoder/disagg_epd_proxy.py" \
414
+ --host "0.0.0.0" \
415
+ --port $PROXY_PORT \
416
+ --encode-servers-urls "http://localhost:$ENCODE_PORT" \
417
+ --prefill-servers-urls "http://localhost:$PREFILL_PORT" \
418
+ --decode-servers-urls "http://localhost:$DECODE_PORT" \
419
+ > $LOG_PATH/1e1p1d_proxy.log 2>&1 &
420
+ PIDS+=($!)
421
+
422
+ # Wait for proxy
423
+ echo "Waiting for proxy..."
424
+ wait_for_server $PROXY_PORT
425
+
426
+ curl http://127.0.0.1:$PROXY_PORT/v1/models
427
+ curl http://127.0.0.1:$PROXY_PORT/health
428
+ echo ""
429
+
430
+ echo "All EPD (1E+1P+1D) services are up!"
431
+
432
+ # Run test in disagg mode
433
+ echo "Running EPD (1E+1P+1D) correctness test..."
434
+
435
+ python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
436
+ --service_url "http://localhost:$PROXY_PORT" \
437
+ --model_name "$MODEL" \
438
+ --mode disagg \
439
+ --baseline_file "$BASELINE_PD_FILE" \
440
+ $MM_FLAG
441
+
442
+ # Cleanup
443
+ echo "✓✓ 1E+1P+1D Correctness Test finished"
444
+ echo "Stopping EPD (1E+1P+1D) instances..."
445
+ for pid in "${PIDS[@]}"; do
446
+ kill $pid 2>/dev/null || true
447
+ done
448
+ sleep 2
449
+ cleanup_instances
450
+}
451
+
452
+# Main execution
453
+echo "================================"
454
+echo "EPD Correctness Test Suite"
455
+echo "Model: $MODEL"
456
+echo "================================"
457
+
458
+# Step 1: Run baseline
459
+run_baseline
460
+
461
+# Step 2: Test 1E + 1PD
462
+run_epd_1e_1pd
463
+
464
+# Step 3: Test baseline 1P + 1D
465
+run_baseline_1p_1d
466
+
467
+# Step 4: Test 1E + 1P + 1D
468
+run_epd_1e_1p_1d
469
+
470
+# Cleanup output file
471
+rm -f "$BASELINE_FILE"
472
+rm -f "$BASELINE_PD_FILE"
473
+
474
+echo "================================"
475
+echo "✓✓ All EPD correctness tests finished!"
476
+echo "================================"
tests/v1/ec_connector/integration/test_epd_correctness.py
new
+305
@@ -0,0 +1,305 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+EPD Correctness Test
5
+
6
+Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same
7
+outputs as a baseline single instance.
8
+
9
+Usage:
10
+ # Baseline mode (saves outputs):
11
+ python test_epd_correctness.py \
12
+ --service_url http://localhost:8000 \
13
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
14
+ --mode baseline \
15
+ --baseline_file .vllm_epd_baseline.txt
16
+
17
+ # Disagg mode (compares outputs):
18
+ python test_epd_correctness.py \
19
+ --service_url http://localhost:8000 \
20
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
21
+ --mode disagg \
22
+ --baseline_file .vllm_epd_baseline.txt
23
+"""
24
+
25
+import argparse
26
+import json
27
+import os
28
+import time
29
+
30
+import openai
31
+import requests
32
+
33
+from vllm.assets.image import ImageAsset
34
+from vllm.multimodal.utils import encode_image_base64
35
+
36
+MAX_OUTPUT_LEN = 256
37
+
38
+# Sample prompts with multimodal content
39
+image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720))
40
+image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720))
41
+
42
+image_local_path = f"{os.path.dirname(os.path.abspath(__file__))}/hato.jpg"
43
+
44
+SAMPLE_PROMPTS_MM: list[dict] = [
45
+ {
46
+ "messages": [
47
+ {
48
+ "role": "user",
49
+ "content": [
50
+ {
51
+ "type": "image_url",
52
+ "image_url": {
53
+ "url": f"data:image;base64,{encode_image_base64(image_1)}"
54
+ },
55
+ },
56
+ {"type": "text", "text": "What's in this image?"},
57
+ ],
58
+ }
59
+ ],
60
+ "description": "Single image query",
61
+ },
62
+ {
63
+ "messages": [
64
+ {
65
+ "role": "user",
66
+ "content": [
67
+ {
68
+ "type": "image_url",
69
+ "image_url": {
70
+ "url": f"data:image;base64,{encode_image_base64(image_2)}"
71
+ },
72
+ },
73
+ {
74
+ "type": "image_url",
75
+ "image_url": {"url": f"file://{image_local_path}"},
76
+ },
77
+ {"type": "text", "text": "Describe these 2 images in detail."},
78
+ ],
79
+ }
80
+ ],
81
+ "description": "2 images with detailed query",
82
+ },
83
+]
84
+
85
+# Text-only prompts for mixed testing
86
+SAMPLE_PROMPTS_TEXT: list[dict] = [
87
+ {
88
+ "messages": [{"role": "user", "content": "What is the capital of France?"}],
89
+ "description": "Simple text-only query",
90
+ },
91
+ {
92
+ "messages": [
93
+ {"role": "user", "content": "Explain quantum computing in simple terms."}
94
+ ],
95
+ "description": "Text-only explanation request",
96
+ },
97
+]
98
+
99
+
100
+def check_vllm_server(url: str, timeout=5, retries=10) -> bool:
101
+ """Check if the vLLM server is ready.
102
+
103
+ Args:
104
+ url: The URL to check (usually /health or /healthcheck endpoint)
105
+ timeout: Timeout in seconds for each request
106
+ retries: Number of retries if the server is not ready
107
+
108
+ Returns:
109
+ True if the server is ready, False otherwise
110
+ """
111
+ for attempt in range(retries):
112
+ try:
113
+ response = requests.get(url, timeout=timeout)
114
+ if response.status_code == 200:
115
+ print(f"Server is ready at {url}")
116
+ return True
117
+ else:
118
+ print(
119
+ f"Attempt {attempt + 1}/{retries}: Server returned "
120
+ f"status code {response.status_code}"
121
+ )
122
+ except requests.exceptions.RequestException as e:
123
+ print(f"Attempt {attempt + 1}/{retries}: Error connecting: {e}")
124
+ time.sleep(2) # Wait before retrying
125
+ return False
126
+
127
+
128
+def run_chat_completion(
129
+ base_url: str,
130
+ model_name: str,
131
+ messages: list,
132
+ max_tokens: int = MAX_OUTPUT_LEN,
133
+) -> str:
134
+ """Run a chat completion request.
135
+
136
+ Args:
137
+ base_url: Base URL of the vLLM server
138
+ model_name: Name of the model
139
+ messages: Messages for chat completion
140
+ max_tokens: Maximum tokens to generate
141
+
142
+ Returns:
143
+ Generated text content
144
+ """
145
+ client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
146
+
147
+ completion = client.chat.completions.create(
148
+ model=model_name,
149
+ messages=messages,
150
+ max_tokens=max_tokens,
151
+ temperature=0.0,
152
+ seed=42,
153
+ )
154
+
155
+ return completion.choices[0].message.content
156
+
157
+
158
+def main():
159
+ """Main test function."""
160
+ parser = argparse.ArgumentParser(
161
+ description="EPD correctness test - compare disagg vs baseline"
162
+ )
163
+
164
+ parser.add_argument(
165
+ "--service_url",
166
+ type=str,
167
+ required=True,
168
+ help="The vLLM service URL (e.g., http://localhost:8000)",
169
+ )
170
+
171
+ parser.add_argument(
172
+ "--model_name",
173
+ type=str,
174
+ required=True,
175
+ help="Model name",
176
+ )
177
+
178
+ parser.add_argument(
179
+ "--mode",
180
+ type=str,
181
+ default="baseline",
182
+ choices=["baseline", "baseline_pd", "disagg"],
183
+ help="Mode: baseline/baseline_pd (saves outputs) or disagg (compares outputs)",
184
+ )
185
+
186
+ parser.add_argument(
187
+ "--baseline_file",
188
+ type=str,
189
+ default=".vllm_epd_baseline.txt",
190
+ help="File to save/load baseline outputs",
191
+ )
192
+
193
+ parser.add_argument(
194
+ "--use_mm_prompts",
195
+ action="store_true",
196
+ help="Use multimodal prompts (default: use text-only for quick testing)",
197
+ )
198
+
199
+ args = parser.parse_args()
200
+
201
+ print(f"Service URL: {args.service_url}")
202
+ print(f"Model: {args.model_name}")
203
+ print(f"Mode: {args.mode}")
204
+ print(f"Output file: {args.baseline_file}")
205
+ print(f"Use MM prompts: {args.use_mm_prompts}")
206
+
207
+ # Determine health check endpoint
208
+ if args.mode == "baseline":
209
+ health_check_url = f"{args.service_url}/health"
210
+ elif args.mode == "baseline_pd":
211
+ # Nixl toy proxy use /healthcheck
212
+ health_check_url = f"{args.service_url}/healthcheck"
213
+ else:
214
+ # Disagg EPD proxy uses /health
215
+ health_check_url = f"{args.service_url}/health"
216
+ if not os.path.exists(args.baseline_file):
217
+ raise ValueError(
218
+ f"In disagg mode, the output file {args.baseline_file} from "
219
+ "baseline does not exist. Run baseline mode first."
220
+ )
221
+
222
+ # Check if server is ready
223
+ if not check_vllm_server(health_check_url):
224
+ raise RuntimeError(f"vLLM server at {args.service_url} is not ready!")
225
+
226
+ # Select prompts to use
227
+ if args.use_mm_prompts:
228
+ test_prompts = SAMPLE_PROMPTS_MM
229
+ print("Using multimodal prompts")
230
+ else:
231
+ test_prompts = SAMPLE_PROMPTS_TEXT
232
+ print("Using text-only prompts for quick testing")
233
+
234
+ # Run completions
235
+ service_url = f"{args.service_url}/v1"
236
+ output_strs = {}
237
+
238
+ for i, prompt_data in enumerate(test_prompts):
239
+ print(
240
+ f"\nRunning prompt {i + 1}/{len(test_prompts)}: {
241
+ prompt_data['description']
242
+ }"
243
+ )
244
+
245
+ output_str = run_chat_completion(
246
+ base_url=service_url,
247
+ model_name=args.model_name,
248
+ messages=prompt_data["messages"],
249
+ max_tokens=MAX_OUTPUT_LEN,
250
+ )
251
+
252
+ # Use description as key for comparison
253
+ key = prompt_data["description"]
254
+ output_strs[key] = output_str
255
+ print(f"Output: {output_str}")
256
+
257
+ if args.mode in ("baseline", "baseline_pd"):
258
+ # Baseline mode: Save outputs
259
+ print(f"\nSaving baseline outputs to {args.baseline_file}")
260
+ try:
261
+ with open(args.baseline_file, "w") as json_file:
262
+ json.dump(output_strs, json_file, indent=4)
263
+ print("✅ Baseline outputs saved successfully")
264
+ except OSError as e:
265
+ print(f"Error writing to file: {e}")
266
+ raise
267
+ else:
268
+ # Disagg mode: Load and compare outputs
269
+ print(f"\nLoading baseline outputs from {args.baseline_file}")
270
+ baseline_outputs = None
271
+ try:
272
+ with open(args.baseline_file) as json_file:
273
+ baseline_outputs = json.load(json_file)
274
+ except OSError as e:
275
+ print(f"Error reading from file: {e}")
276
+ raise
277
+
278
+ # Verify outputs match
279
+ print("\nComparing disagg outputs with baseline...")
280
+ assert isinstance(baseline_outputs, dict), "Baseline outputs should be a dict"
281
+ assert len(baseline_outputs) == len(output_strs), (
282
+ f"Length mismatch: baseline has {len(baseline_outputs)}, "
283
+ f"disagg has {len(output_strs)}"
284
+ )
285
+
286
+ all_match = True
287
+ for key, baseline_output in baseline_outputs.items():
288
+ assert key in output_strs, f"{key} not in disagg outputs"
289
+
290
+ disagg_output = output_strs[key]
291
+ if baseline_output == disagg_output:
292
+ print(f"✅ {key}: MATCH")
293
+ else:
294
+ print(f"❌ {key}: MISMATCH")
295
+ print(f" Baseline: {baseline_output}")
296
+ print(f" Disagg: {disagg_output}")
297
+ all_match = False
298
+
299
+ assert all_match, "❌❌Disagg outputs do not match baseline!❌❌"
300
+ if all_match:
301
+ print("\n✅ All outputs match! Test PASSED")
302
+
303
+
304
+if __name__ == "__main__":
305
+ main()
tests/v1/ec_connector/unit/test_ec_shared_storage_connector.py
new
+609
@@ -0,0 +1,609 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for ECSharedStorageConnector.
5
+"""
6
+
7
+import os
8
+from unittest.mock import Mock, patch
9
+
10
+import pytest
11
+import safetensors
12
+import torch
13
+
14
+from vllm.config import VllmConfig
15
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
16
+from vllm.distributed.ec_transfer.ec_connector.shared_storage_connector import (
17
+ ECSharedStorageConnector,
18
+ ECSharedStorageConnectorMetadata,
19
+ MMMeta,
20
+)
21
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
22
+from vllm.v1.core.sched.output import SchedulerOutput
23
+
24
+
25
+# ------------------ Mock Classes ------------------ #
26
+class MockRequest:
27
+ def __init__(self, request_id, mm_hashes: list[str], token_counts: list[int]):
28
+ assert len(mm_hashes) == len(token_counts)
29
+ self.request_id = request_id
30
+ self._token_counts = token_counts
31
+ self.mm_features = []
32
+ for i, mm_hash in enumerate(mm_hashes):
33
+ feature = MultiModalFeatureSpec(
34
+ data=None,
35
+ modality="image",
36
+ identifier=mm_hash,
37
+ mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]),
38
+ )
39
+ self.mm_features.append(feature)
40
+
41
+ def get_num_encoder_tokens(self, input_id: int) -> int:
42
+ assert input_id < len(self._token_counts)
43
+ return self._token_counts[input_id]
44
+
45
+
46
+@pytest.fixture
47
+def temp_storage(tmp_path):
48
+ """Fixture providing temporary storage path."""
49
+ return str(tmp_path)
50
+
51
+
52
+@pytest.fixture
53
+def mock_vllm_config_producer(temp_storage):
54
+ """Fixture providing mock VllmConfig for producer role."""
55
+ config = Mock(spec=VllmConfig)
56
+ config.ec_transfer_config = Mock()
57
+ config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
58
+ config.ec_transfer_config.is_ec_producer = True
59
+ return config
60
+
61
+
62
+@pytest.fixture
63
+def mock_vllm_config_consumer(temp_storage):
64
+ """Fixture providing mock VllmConfig for consumer role."""
65
+ config = Mock(spec=VllmConfig)
66
+ config.ec_transfer_config = Mock()
67
+ config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
68
+ config.ec_transfer_config.is_ec_producer = False
69
+ return config
70
+
71
+
72
+@pytest.fixture
73
+def mock_request_with_3_mm():
74
+ """Fixture providing mock Request with 3 multimodal items."""
75
+ request_id = "test_req_123"
76
+ mm_hashes = ["img_hash_1", "img_hash_2", "img_hash_3"]
77
+ token_counts = [100, 150, 200]
78
+
79
+ request = MockRequest(request_id, mm_hashes, token_counts)
80
+ return request
81
+
82
+
83
+# ------------------ Unit Tests ------------------ #
84
+class TestECSharedStorageConnectorBasics:
85
+ """Test basic EC connector functionality."""
86
+
87
+ def test_initialization_producer(self, mock_vllm_config_producer, temp_storage):
88
+ """Test connector initializes correctly as producer."""
89
+ connector = ECSharedStorageConnector(
90
+ vllm_config=mock_vllm_config_producer,
91
+ role=ECConnectorRole.SCHEDULER,
92
+ )
93
+
94
+ assert connector.role == ECConnectorRole.SCHEDULER
95
+ assert connector.is_producer
96
+ assert connector._storage_path == temp_storage
97
+ assert connector._mm_datas_need_loads == {}
98
+
99
+ def test_initialization_consumer(self, mock_vllm_config_consumer, temp_storage):
100
+ """Test connector initializes correctly as consumer."""
101
+ connector = ECSharedStorageConnector(
102
+ vllm_config=mock_vllm_config_consumer,
103
+ role=ECConnectorRole.WORKER,
104
+ )
105
+
106
+ assert connector.role == ECConnectorRole.WORKER
107
+ assert not connector.is_producer
108
+ assert connector._storage_path == temp_storage
109
+
110
+ def test_role_assignment(self, mock_vllm_config_producer):
111
+ """Test role is correctly assigned."""
112
+ scheduler_connector = ECSharedStorageConnector(
113
+ vllm_config=mock_vllm_config_producer,
114
+ role=ECConnectorRole.SCHEDULER,
115
+ )
116
+ worker_connector = ECSharedStorageConnector(
117
+ vllm_config=mock_vllm_config_producer,
118
+ role=ECConnectorRole.WORKER,
119
+ )
120
+
121
+ assert scheduler_connector.role == ECConnectorRole.SCHEDULER
122
+ assert worker_connector.role == ECConnectorRole.WORKER
123
+
124
+
125
+class TestCacheExistence:
126
+ """Test cache existence checking using has_caches() API."""
127
+
128
+ def test_has_caches_all_exist_3_items(
129
+ self,
130
+ mock_vllm_config_producer,
131
+ mock_vllm_config_consumer,
132
+ mock_request_with_3_mm,
133
+ ):
134
+ """Test has_caches returns True when all 3 caches exist."""
135
+ # Test for producer first
136
+ producer = ECSharedStorageConnector(
137
+ vllm_config=mock_vllm_config_producer,
138
+ role=ECConnectorRole.SCHEDULER,
139
+ )
140
+
141
+ # Create cache files using save_caches (proper way)
142
+ encoder_cache: dict[str, torch.Tensor] = {}
143
+
144
+ for mm_feature in mock_request_with_3_mm.mm_features:
145
+ mm_hash = mm_feature.identifier
146
+ encoder_cache[mm_hash] = torch.randn(10, 768)
147
+ producer.save_caches(encoder_cache, mm_hash)
148
+
149
+ # Test using has_caches API
150
+ producer_result = producer.has_caches(mock_request_with_3_mm)
151
+
152
+ # Assert
153
+ assert len(producer_result) == 3
154
+ assert all(producer_result), f"Expected all True, got {producer_result}"
155
+
156
+ # Also test consumer can check if cache exists
157
+ consumer = ECSharedStorageConnector(
158
+ vllm_config=mock_vllm_config_consumer,
159
+ role=ECConnectorRole.SCHEDULER,
160
+ )
161
+
162
+ # Test using has_caches API
163
+ consumer_result = consumer.has_caches(mock_request_with_3_mm)
164
+
165
+ # Assert
166
+ assert len(consumer_result) == 3
167
+ assert all(consumer_result), f"Expected all True, got {consumer_result}"
168
+
169
+ def test_has_caches_none_exist(
170
+ self, mock_vllm_config_producer, mock_request_with_3_mm
171
+ ):
172
+ """Test has_caches returns False when no caches exist."""
173
+ connector = ECSharedStorageConnector(
174
+ vllm_config=mock_vllm_config_producer,
175
+ role=ECConnectorRole.SCHEDULER,
176
+ )
177
+
178
+ # Test without creating any files
179
+ result = connector.has_caches(mock_request_with_3_mm)
180
+
181
+ # Assert
182
+ assert len(result) == 3
183
+ assert not any(result), f"Expected all False, got {result}"
184
+
185
+ def test_has_caches_partial_exist(
186
+ self, mock_vllm_config_producer, mock_request_with_3_mm
187
+ ):
188
+ """Test has_caches with some caches existing (1 of 3)."""
189
+ connector = ECSharedStorageConnector(
190
+ vllm_config=mock_vllm_config_producer,
191
+ role=ECConnectorRole.SCHEDULER,
192
+ )
193
+
194
+ # Create only the second cache file
195
+ mm_hash_second = mock_request_with_3_mm.mm_features[1].identifier
196
+ encoder_cache = {mm_hash_second: torch.randn(10, 768)}
197
+ connector.save_caches(encoder_cache, mm_hash_second)
198
+
199
+ # Test
200
+ result = connector.has_caches(mock_request_with_3_mm)
201
+
202
+ # Assert
203
+ assert len(result) == 3
204
+ assert not result[0] # First doesn't exist
205
+ assert result[1] # Second exists
206
+ assert not result[2] # Third doesn't exist
207
+
208
+
209
+class TestStateManagement:
210
+ """Test connector state management."""
211
+
212
+ def test_update_state_after_alloc_3_items(
213
+ self, mock_vllm_config_producer, mock_request_with_3_mm
214
+ ):
215
+ """Test state update after allocation for 3 MM items."""
216
+ connector = ECSharedStorageConnector(
217
+ vllm_config=mock_vllm_config_producer,
218
+ role=ECConnectorRole.SCHEDULER,
219
+ )
220
+
221
+ # Initial state should be empty
222
+ assert len(connector._mm_datas_need_loads) == 0
223
+
224
+ # Update state for all 3 items
225
+ for i in range(3):
226
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
227
+
228
+ # Check state updated for all 3
229
+ assert len(connector._mm_datas_need_loads) == 3
230
+ assert "img_hash_1" in connector._mm_datas_need_loads
231
+ assert "img_hash_2" in connector._mm_datas_need_loads
232
+ assert "img_hash_3" in connector._mm_datas_need_loads
233
+ assert connector._mm_datas_need_loads["img_hash_1"] == 100
234
+ assert connector._mm_datas_need_loads["img_hash_2"] == 150
235
+ assert connector._mm_datas_need_loads["img_hash_3"] == 200
236
+
237
+ def test_build_connector_meta_3_items(
238
+ self, mock_vllm_config_producer, mock_request_with_3_mm
239
+ ):
240
+ """Test metadata building for 3 MM items."""
241
+ connector = ECSharedStorageConnector(
242
+ vllm_config=mock_vllm_config_producer,
243
+ role=ECConnectorRole.SCHEDULER,
244
+ )
245
+
246
+ # Setup state for all 3 items
247
+ for i in range(3):
248
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
249
+
250
+ # Build metadata
251
+ scheduler_output = Mock(spec=SchedulerOutput)
252
+ metadata = connector.build_connector_meta(scheduler_output)
253
+
254
+ # Assert
255
+ assert isinstance(metadata, ECSharedStorageConnectorMetadata)
256
+ assert len(metadata.mm_datas) == 3
257
+ assert metadata.mm_datas[0].mm_hash == "img_hash_1"
258
+ assert metadata.mm_datas[0].num_token == 100
259
+ assert metadata.mm_datas[1].mm_hash == "img_hash_2"
260
+ assert metadata.mm_datas[1].num_token == 150
261
+ assert metadata.mm_datas[2].mm_hash == "img_hash_3"
262
+ assert metadata.mm_datas[2].num_token == 200
263
+
264
+ # State should be cleared after building
265
+ assert len(connector._mm_datas_need_loads) == 0
266
+
267
+ def test_build_connector_meta_empty(self, mock_vllm_config_producer):
268
+ """Test metadata building with empty state."""
269
+ connector = ECSharedStorageConnector(
270
+ vllm_config=mock_vllm_config_producer,
271
+ role=ECConnectorRole.SCHEDULER,
272
+ )
273
+
274
+ scheduler_output = Mock(spec=SchedulerOutput)
275
+ metadata = connector.build_connector_meta(scheduler_output)
276
+
277
+ assert isinstance(metadata, ECSharedStorageConnectorMetadata)
278
+ assert len(metadata.mm_datas) == 0
279
+
280
+ def test_state_cleared_after_metadata_build(
281
+ self, mock_vllm_config_producer, mock_request_with_3_mm
282
+ ):
283
+ """Test that state is properly cleared after building metadata."""
284
+ connector = ECSharedStorageConnector(
285
+ vllm_config=mock_vllm_config_producer,
286
+ role=ECConnectorRole.SCHEDULER,
287
+ )
288
+
289
+ # Add state
290
+ for i in range(3):
291
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
292
+ assert len(connector._mm_datas_need_loads) == 3
293
+
294
+ # Build metadata (should clear state)
295
+ scheduler_output = Mock(spec=SchedulerOutput)
296
+ connector.build_connector_meta(scheduler_output)
297
+
298
+ # State should be empty
299
+ assert len(connector._mm_datas_need_loads) == 0
300
+
301
+ # Build again should return empty metadata
302
+ metadata2 = connector.build_connector_meta(scheduler_output)
303
+ assert len(metadata2.mm_datas) == 0
304
+
305
+
306
+class TestCacheSaving:
307
+ """Test encoder cache saving (producer only)."""
308
+
309
+ def test_save_caches_producer_3_items(
310
+ self, mock_vllm_config_producer, mock_request_with_3_mm, temp_storage
311
+ ):
312
+ """Test cache saving as producer for 3 different MM items."""
313
+ connector = ECSharedStorageConnector(
314
+ vllm_config=mock_vllm_config_producer,
315
+ role=ECConnectorRole.WORKER,
316
+ )
317
+
318
+ # Create and save 3 different caches
319
+ mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
320
+ encoder_cache: dict[str, torch.Tensor] = {}
321
+
322
+ for mm_hash in mm_hashes:
323
+ encoder_cache[mm_hash] = torch.randn(10, 768)
324
+ connector.save_caches(encoder_cache, mm_hash)
325
+
326
+ # Verify all files exist using has_caches
327
+ result = connector.has_caches(mock_request_with_3_mm)
328
+ assert all(result), f"Not all caches were saved: {result}"
329
+
330
+ # Verify each file's content
331
+ for mm_hash in mm_hashes:
332
+ filename = connector._generate_filename_debug(mm_hash)
333
+ loaded = safetensors.torch.load_file(filename)
334
+ assert "ec_cache" in loaded
335
+ assert torch.allclose(loaded["ec_cache"], encoder_cache[mm_hash].cpu())
336
+
337
+ def test_save_caches_consumer_skips(self, mock_vllm_config_consumer):
338
+ """Test cache saving is skipped for consumer."""
339
+ connector = ECSharedStorageConnector(
340
+ vllm_config=mock_vllm_config_consumer,
341
+ role=ECConnectorRole.WORKER,
342
+ )
343
+
344
+ mm_hash = "test_hash_consumer"
345
+ encoder_cache = {mm_hash: torch.randn(10, 768)}
346
+
347
+ # Save should not raise but also not create file
348
+ connector.save_caches(encoder_cache, mm_hash)
349
+
350
+ # Verify file doesn't exist using has_caches
351
+ mock_request = MockRequest("req_consumer", [mm_hash], [10])
352
+ result = connector.has_caches(mock_request)
353
+ assert not result[0], "Consumer should not save caches"
354
+
355
+
356
+class TestCacheLoading:
357
+ """Test encoder cache loading (consumer)."""
358
+
359
+ @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
360
+ def test_start_load_caches_consumer_3_items(
361
+ self,
362
+ mock_vllm_config_producer,
363
+ mock_vllm_config_consumer,
364
+ mock_request_with_3_mm,
365
+ temp_storage,
366
+ ):
367
+ """Test consumer loads 3 caches from storage."""
368
+ # First, create producer to save caches
369
+ producer = ECSharedStorageConnector(
370
+ vllm_config=mock_vllm_config_producer,
371
+ role=ECConnectorRole.WORKER,
372
+ )
373
+
374
+ # Producer saves 3 caches
375
+ mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
376
+ saved_caches = {}
377
+ for mm_hash in mm_hashes:
378
+ saved_caches[mm_hash] = torch.randn(10, 768)
379
+ producer.save_caches(saved_caches, mm_hash)
380
+
381
+ # Now consumer loads
382
+ consumer = ECSharedStorageConnector(
383
+ vllm_config=mock_vllm_config_consumer,
384
+ role=ECConnectorRole.WORKER,
385
+ )
386
+
387
+ # Setup metadata for all 3
388
+ metadata = ECSharedStorageConnectorMetadata()
389
+ for mm_hash in mm_hashes:
390
+ metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
391
+ consumer.bind_connector_metadata(metadata)
392
+
393
+ # Load
394
+ encoder_cache: dict[str, torch.Tensor] = {}
395
+ consumer.start_load_caches(encoder_cache=encoder_cache)
396
+
397
+ # Verify all 3 loaded
398
+ assert len(encoder_cache) == 3
399
+ for mm_hash in mm_hashes:
400
+ assert mm_hash in encoder_cache, f"{mm_hash} missing in encoder_cache"
401
+ assert encoder_cache[mm_hash].is_cuda, (
402
+ f"{mm_hash} cache is in {encoder_cache[mm_hash].device}"
403
+ )
404
+ assert torch.allclose(
405
+ encoder_cache[mm_hash].cpu(), saved_caches[mm_hash]
406
+ ), f"{mm_hash} cache saved and loaded tesnor are not the same"
407
+
408
+ def test_start_load_caches_skip_existing(
409
+ self, mock_vllm_config_producer, mock_vllm_config_consumer, temp_storage
410
+ ):
411
+ """Test cache loading skips already cached items."""
412
+ # Setup: producer saves cache
413
+ producer = ECSharedStorageConnector(
414
+ vllm_config=mock_vllm_config_producer,
415
+ role=ECConnectorRole.WORKER,
416
+ )
417
+
418
+ mm_hash = "existing_hash"
419
+ saved_cache = torch.randn(10, 768)
420
+ producer.save_caches({mm_hash: saved_cache}, mm_hash)
421
+
422
+ # Consumer setup
423
+ consumer = ECSharedStorageConnector(
424
+ vllm_config=mock_vllm_config_consumer,
425
+ role=ECConnectorRole.WORKER,
426
+ )
427
+
428
+ metadata = ECSharedStorageConnectorMetadata()
429
+ metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
430
+ consumer.bind_connector_metadata(metadata)
431
+
432
+ # Pre-populate encoder_cache with different value
433
+ existing_cache = torch.randn(5, 512)
434
+ encoder_cache = {mm_hash: existing_cache}
435
+
436
+ # Load (should skip since already exists)
437
+ with patch("safetensors.torch.load_file") as mock_load:
438
+ consumer.start_load_caches(encoder_cache=encoder_cache)
439
+ # Should not call load_file since cache exists
440
+ mock_load.assert_not_called()
441
+
442
+ # Verify original cache unchanged
443
+ assert torch.equal(encoder_cache[mm_hash], existing_cache)
444
+
445
+ def test_start_load_caches_empty_metadata(self, mock_vllm_config_consumer):
446
+ """Test loading with empty metadata does nothing."""
447
+ consumer = ECSharedStorageConnector(
448
+ vllm_config=mock_vllm_config_consumer,
449
+ role=ECConnectorRole.WORKER,
450
+ )
451
+
452
+ # Setup empty metadata
453
+ metadata = ECSharedStorageConnectorMetadata()
454
+ consumer.bind_connector_metadata(metadata)
455
+
456
+ # Load (should not raise)
457
+ encoder_cache: dict[str, torch.Tensor] = {}
458
+ consumer.start_load_caches(encoder_cache=encoder_cache)
459
+
460
+ # Cache should remain empty
461
+ assert len(encoder_cache) == 0
462
+
463
+
464
+class TestFilenameGeneration:
465
+ """Test filename and path generation."""
466
+
467
+ def test_generate_foldername(self, mock_vllm_config_producer, temp_storage):
468
+ """Test folder name generation."""
469
+ connector = ECSharedStorageConnector(
470
+ vllm_config=mock_vllm_config_producer,
471
+ role=ECConnectorRole.WORKER,
472
+ )
473
+
474
+ mm_hash = "test_folder_hash"
475
+ folder = connector._generate_foldername_debug(mm_hash)
476
+
477
+ assert folder == os.path.join(temp_storage, mm_hash)
478
+ assert os.path.isdir(folder) # Should be created
479
+
480
+ def test_generate_filename(self, mock_vllm_config_producer, temp_storage):
481
+ """Test filename generation."""
482
+ connector = ECSharedStorageConnector(
483
+ vllm_config=mock_vllm_config_producer,
484
+ role=ECConnectorRole.WORKER,
485
+ )
486
+
487
+ mm_hash = "test_file_hash"
488
+ filename = connector._generate_filename_debug(mm_hash)
489
+
490
+ expected = os.path.join(temp_storage, mm_hash, "encoder_cache.safetensors")
491
+ assert filename == expected
492
+ assert os.path.isdir(os.path.dirname(filename)) # Folder created
493
+
494
+ def test_generate_filename_consistency(self, mock_vllm_config_producer):
495
+ """Test filename generation is consistent."""
496
+ connector = ECSharedStorageConnector(
497
+ vllm_config=mock_vllm_config_producer,
498
+ role=ECConnectorRole.WORKER,
499
+ )
500
+
501
+ mm_hash = "consistency_hash"
502
+ filename1 = connector._generate_filename_debug(mm_hash)
503
+ filename2 = connector._generate_filename_debug(mm_hash)
504
+
505
+ assert filename1 == filename2
506
+
507
+
508
+class TestMetadataBindingLifecycle:
509
+ """Test metadata binding and clearing lifecycle."""
510
+
511
+ def test_bind_connector_metadata(self, mock_vllm_config_consumer):
512
+ """Test binding connector metadata."""
513
+ connector = ECSharedStorageConnector(
514
+ vllm_config=mock_vllm_config_consumer,
515
+ role=ECConnectorRole.WORKER,
516
+ )
517
+
518
+ metadata = ECSharedStorageConnectorMetadata()
519
+ metadata.add_mm_data(MMMeta.make_meta("hash_1", 100))
520
+
521
+ connector.bind_connector_metadata(metadata)
522
+
523
+ assert connector._connector_metadata is metadata
524
+
525
+ def test_clear_connector_metadata(self, mock_vllm_config_consumer):
526
+ """Test clearing connector metadata."""
527
+ connector = ECSharedStorageConnector(
528
+ vllm_config=mock_vllm_config_consumer,
529
+ role=ECConnectorRole.WORKER,
530
+ )
531
+
532
+ metadata = ECSharedStorageConnectorMetadata()
533
+ connector.bind_connector_metadata(metadata)
534
+
535
+ connector.clear_connector_metadata()
536
+
537
+ assert connector._connector_metadata is None
538
+
539
+ def test_get_connector_metadata(self, mock_vllm_config_consumer):
540
+ """Test getting connector metadata."""
541
+ connector = ECSharedStorageConnector(
542
+ vllm_config=mock_vllm_config_consumer,
543
+ role=ECConnectorRole.WORKER,
544
+ )
545
+
546
+ metadata = ECSharedStorageConnectorMetadata()
547
+ connector.bind_connector_metadata(metadata)
548
+
549
+ retrieved = connector._get_connector_metadata()
550
+
551
+ assert retrieved is metadata
552
+
553
+ def test_get_connector_metadata_not_set(self, mock_vllm_config_consumer):
554
+ """Test getting metadata when not set raises."""
555
+ connector = ECSharedStorageConnector(
556
+ vllm_config=mock_vllm_config_consumer,
557
+ role=ECConnectorRole.WORKER,
558
+ )
559
+
560
+ with pytest.raises(AssertionError):
561
+ connector._get_connector_metadata()
562
+
563
+
564
+class TestEdgeCases:
565
+ """Test edge cases and error handling."""
566
+
567
+ def test_save_empty_cache(self, mock_vllm_config_producer):
568
+ """Test saving empty tensor."""
569
+ connector = ECSharedStorageConnector(
570
+ vllm_config=mock_vllm_config_producer,
571
+ role=ECConnectorRole.WORKER,
572
+ )
573
+
574
+ mm_hash = "empty_hash"
575
+ encoder_cache = {mm_hash: torch.empty(0)}
576
+
577
+ # Should not raise
578
+ connector.save_caches(encoder_cache, mm_hash)
579
+
580
+ def test_load_nonexistent_cache(self, mock_vllm_config_consumer):
581
+ """Test loading cache that doesn't exist raises error."""
582
+ connector = ECSharedStorageConnector(
583
+ vllm_config=mock_vllm_config_consumer,
584
+ role=ECConnectorRole.WORKER,
585
+ )
586
+
587
+ metadata = ECSharedStorageConnectorMetadata()
588
+ metadata.add_mm_data(MMMeta.make_meta("nonexistent_hash", 100))
589
+ connector.bind_connector_metadata(metadata)
590
+
591
+ encoder_cache: dict[str, torch.Tensor] = {}
592
+
593
+ # Should raise FileNotFoundError
594
+ with pytest.raises(FileNotFoundError):
595
+ connector.start_load_caches(encoder_cache=encoder_cache)
596
+
597
+ def test_has_caches_empty_request(self, mock_vllm_config_producer):
598
+ """Test has_caches with request that has no MM data."""
599
+ connector = ECSharedStorageConnector(
600
+ vllm_config=mock_vllm_config_producer,
601
+ role=ECConnectorRole.SCHEDULER,
602
+ )
603
+
604
+ mock_request = MockRequest("req_empty", [], [])
605
+
606
+ result = connector.has_caches(mock_request)
607
+
608
+ assert len(result) == 0
609
+ assert result == []
tests/v1/engine/test_engine_core.py
+146
@@ -10,6 +10,14 @@ import pytest
10
from transformers import AutoTokenizer
11
12
from vllm import SamplingParams
13
+from vllm.config import (
14
+ CacheConfig,
15
+ ECTransferConfig,
16
+ KVTransferConfig,
17
+ ModelConfig,
18
+ SchedulerConfig,
19
+ VllmConfig,
20
+)
21
from vllm.engine.arg_utils import EngineArgs
22
from vllm.platforms import current_platform
23
from vllm.utils.torch_utils import set_default_torch_num_threads
@@ -450,3 +458,141 @@ def test_engine_core_invalid_request_id_type():
458
engine_core.add_request(*engine_core.preprocess_add_request(valid_request))
459
assert len(engine_core.scheduler.waiting) == 1
460
assert len(engine_core.scheduler.running) == 0
461
+
462
+
463
+@create_new_process_for_each_test()
464
+@pytest.mark.parametrize(
465
+ ("ec_role", "gpu_memory_utilization", "enable_prefix_caching"),
466
+ [
467
+ ("ec_producer", 0.01, False),
468
+ # NOTE: ec_producer never allows prefix caching
469
+ ("ec_consumer", 0.7, True),
470
+ ("ec_consumer", 0.7, False),
471
+ ],
472
+)
473
+@pytest.mark.parametrize("use_kv_connector", [False, True])
474
+def test_encoder_instance_zero_kv_cache(
475
+ ec_role: str,
476
+ gpu_memory_utilization: float,
477
+ enable_prefix_caching: bool,
478
+ use_kv_connector: bool,
479
+):
480
+ """EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests
481
+
482
+ This test verifies encoder-only instance initializes with 0 KV cache blocks.
483
+ Under EPD disagg mode, Encoder instances (EC producer role) only execute
484
+ vision encoder, so they don't need KV cache for text generation.
485
+ """
486
+ # Form vllm config
487
+ scheduler_config = SchedulerConfig(
488
+ max_num_seqs=10,
489
+ max_num_batched_tokens=512,
490
+ max_model_len=512,
491
+ disable_hybrid_kv_cache_manager=True,
492
+ )
493
+ model_config = ModelConfig(
494
+ model="llava-hf/llava-1.5-7b-hf", # Multimodal model
495
+ enforce_eager=True,
496
+ trust_remote_code=True,
497
+ dtype="float16",
498
+ seed=42,
499
+ )
500
+ cache_config = CacheConfig(
501
+ block_size=16,
502
+ gpu_memory_utilization=gpu_memory_utilization,
503
+ swap_space=0,
504
+ cache_dtype="auto",
505
+ enable_prefix_caching=enable_prefix_caching,
506
+ )
507
+ kv_transfer_config = (
508
+ KVTransferConfig(
509
+ kv_connector="SharedStorageConnector",
510
+ kv_role="kv_both",
511
+ kv_connector_extra_config={"shared_storage_path": "local_storage"},
512
+ )
513
+ if use_kv_connector
514
+ else None
515
+ )
516
+ ec_transfer_config = ECTransferConfig(
517
+ ec_connector="ECSharedStorageConnector",
518
+ ec_role=ec_role,
519
+ ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test_encoder"},
520
+ )
521
+
522
+ vllm_config = VllmConfig(
523
+ model_config=model_config,
524
+ cache_config=cache_config,
525
+ scheduler_config=scheduler_config,
526
+ kv_transfer_config=kv_transfer_config,
527
+ ec_transfer_config=ec_transfer_config,
528
+ )
529
+
530
+ executor_class = Executor.get_class(vllm_config)
531
+ print(f"executor_class: {executor_class}")
532
+
533
+ with set_default_torch_num_threads(1):
534
+ engine_core = EngineCore(
535
+ vllm_config=vllm_config, executor_class=executor_class, log_stats=True
536
+ )
537
+
538
+ # Check encoder cache manager exists
539
+ assert engine_core.scheduler.encoder_cache_manager is not None, (
540
+ "encoder_cache_manager should exist"
541
+ )
542
+
543
+ if ec_role == "ec_producer":
544
+ # Check 1: num_blocks should be 0
545
+ # NOTE: num_blocks=1 as BlockPool always needs a null_block.
546
+ kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
547
+ print(f"kv_cache_config: {kv_cache_config}")
548
+ assert kv_cache_config.num_blocks == 1, (
549
+ f"ec_producer should only have 1 KV blocks, "
550
+ f"got {kv_cache_config.num_blocks}"
551
+ )
552
+
553
+ # Check 2: kv_cache_groups should be empty
554
+ assert len(kv_cache_config.kv_cache_groups) == 0, (
555
+ f"ec_producer should have 0 KV cache groups, "
556
+ f"got {len(kv_cache_config.kv_cache_groups)}"
557
+ )
558
+
559
+ # Check 3: kv_cache_tensors should be empty
560
+ assert len(kv_cache_config.kv_cache_tensors) == 0, (
561
+ f"Encoder instance should have 0 KV cache tensors, "
562
+ f"got {len(kv_cache_config.kv_cache_tensors)}"
563
+ )
564
+
565
+ # Check 4: Verify EC connector is initialized and is producer
566
+ assert engine_core.scheduler.ec_connector is not None, (
567
+ "Encoder instance should have EC connector"
568
+ )
569
+ assert engine_core.scheduler.ec_connector.is_producer, (
570
+ "Encoder instance EC connector should be producer"
571
+ )
572
+
573
+ # Check 5: Verify chunked prefill is disabled
574
+ assert not vllm_config.scheduler_config.chunked_prefill_enabled, (
575
+ "Encoder instance should disable chunked prefill (no KV cache)"
576
+ )
577
+
578
+ elif ec_role == "ec_consumer":
579
+ # Check 1: num_blocks should be > 1
580
+ kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
581
+ print(f"kv_cache_config: {kv_cache_config}")
582
+ assert kv_cache_config.num_blocks > 1, (
583
+ f"ec_consumer should have >1 KV blocks, got {kv_cache_config.num_blocks}"
584
+ )
585
+
586
+ # Check 2: kv_cache_groups should NOT be empty
587
+ assert len(kv_cache_config.kv_cache_groups) > 0, (
588
+ f"ec_consumer should have KV cache groups, "
589
+ f"got {len(kv_cache_config.kv_cache_groups)}"
590
+ )
591
+
592
+ # Check 3: Verify EC connector is consumer
593
+ assert engine_core.scheduler.ec_connector is not None, (
594
+ "Consumer instance should have EC connector"
595
+ )
596
+ assert not engine_core.scheduler.ec_connector.is_producer, (
597
+ "Consumer instance EC connector should be consumer"
598
+ )
vllm/config/__init__.py
+3
@@ -9,6 +9,7 @@ from vllm.config.compilation import (
9
PassConfig,
10
)
11
from vllm.config.device import DeviceConfig
12
+from vllm.config.ec_transfer import ECTransferConfig
13
from vllm.config.kv_events import KVEventsConfig
14
from vllm.config.kv_transfer import KVTransferConfig
15
from vllm.config.load import LoadConfig
@@ -54,6 +55,8 @@ __all__ = [
55
"PassConfig",
56
# From vllm.config.device
57
"DeviceConfig",
58
+ # From vllm.config.ec_transfer
59
+ "ECTransferConfig",
60
# From vllm.config.kv_events
61
"KVEventsConfig",
62
# From vllm.config.kv_transfer
vllm/config/ec_transfer.py
new
+110
@@ -0,0 +1,110 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import hashlib
4
+import uuid
5
+from dataclasses import field
6
+from typing import Any, Literal, get_args
7
+
8
+from pydantic.dataclasses import dataclass
9
+
10
+from vllm.config.utils import config
11
+
12
+ECProducer = Literal["ec_producer"]
13
+ECConsumer = Literal["ec_consumer"]
14
+ECRole = Literal[ECProducer, ECConsumer]
15
+
16
+
17
+@config
18
+@dataclass
19
+class ECTransferConfig:
20
+ """Configuration for distributed EC cache transfer."""
21
+
22
+ ec_connector: str | None = None
23
+ """The EC connector for vLLM to transmit EC caches between vLLM instances.
24
+ """
25
+
26
+ engine_id: str | None = None
27
+ """The engine id for EC transfers."""
28
+
29
+ ec_buffer_device: str | None = "cuda"
30
+ """The device used by ec connector to buffer the EC cache.
31
+ Currently only support 'cuda'."""
32
+
33
+ ec_buffer_size: float = 1e9
34
+ """The buffer size for TorchDistributedConnector. Measured in number of
35
+ bytes. Recommended value: 1e9 (about 1GB)."""
36
+
37
+ ec_role: ECRole | None = None
38
+ """Whether this vLLM instance produces, consumes EC cache, or both. Choices
39
+ are 'ec_producer', 'ec_consumer'."""
40
+
41
+ ec_rank: int | None = None
42
+ """The rank of this vLLM instance in the EC cache transfer. Typical value:
43
+ 0 for encoder, 1 for pd instance.
44
+ Currently only 1P1D is supported."""
45
+
46
+ ec_parallel_size: int = 1
47
+ """The number of parallel instances for EC cache transfer. For
48
+ PyNcclConnector, this should be 2."""
49
+
50
+ ec_ip: str = "127.0.0.1"
51
+ """The EC connector ip, used to build distributed connection."""
52
+
53
+ ec_port: int = 14579
54
+ """The EC connector port, used to build distributed connection."""
55
+
56
+ ec_connector_extra_config: dict[str, Any] = field(default_factory=dict)
57
+ """any extra config that the connector may need."""
58
+
59
+ ec_connector_module_path: str | None = None
60
+ """The Python module path to dynamically load the EC connector from.
61
+ Only supported in V1."""
62
+
63
+ def compute_hash(self) -> str:
64
+ """
65
+ WARNING: Whenever a new field is added to this config,
66
+ ensure that it is included in the factors list if
67
+ it affects the computation graph.
68
+
69
+ Provide a hash that uniquely identifies all the configs
70
+ that affect the structure of the computation
71
+ graph from input ids/embeddings to the final hidden states,
72
+ excluding anything before input ids/embeddings and after
73
+ the final hidden states.
74
+ """
75
+ # no factors to consider.
76
+ # this config will not affect the computation graph.
77
+ factors: list[Any] = []
78
+ hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
79
+ return hash_str
80
+
81
+ def __post_init__(self) -> None:
82
+ if self.engine_id is None:
83
+ self.engine_id = str(uuid.uuid4())
84
+
85
+ if self.ec_role is not None and self.ec_role not in get_args(ECRole):
86
+ raise ValueError(
87
+ f"Unsupported ec_role: {self.ec_role}. "
88
+ f"Supported roles are {get_args(ECRole)}"
89
+ )
90
+
91
+ if self.ec_connector is not None and self.ec_role is None:
92
+ raise ValueError(
93
+ "Please specify ec_role when ec_connector "
94
+ f"is set, supported roles are {get_args(ECRole)}"
95
+ )
96
+
97
+ @property
98
+ def is_ec_transfer_instance(self) -> bool:
99
+ return self.ec_connector is not None and self.ec_role in get_args(ECRole)
100
+
101
+ @property
102
+ def is_ec_producer(self) -> bool:
103
+ return self.ec_connector is not None and self.ec_role in get_args(ECProducer)
104
+
105
+ @property
106
+ def is_ec_consumer(self) -> bool:
107
+ return self.ec_connector is not None and self.ec_role in get_args(ECConsumer)
108
+
109
+ def get_from_extra_config(self, key, default) -> Any:
110
+ return self.ec_connector_extra_config.get(key, default)
vllm/config/vllm.py
+7
@@ -28,6 +28,7 @@ from vllm.utils import random_uuid
28
from .cache import CacheConfig
29
from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
30
from .device import DeviceConfig
31
+from .ec_transfer import ECTransferConfig
32
from .kv_events import KVEventsConfig
33
from .kv_transfer import KVTransferConfig
34
from .load import LoadConfig
@@ -103,6 +104,8 @@ class VllmConfig:
104
"""The configurations for distributed KV cache transfer."""
105
kv_events_config: KVEventsConfig | None = None
106
"""The configurations for event publishing."""
107
+ ec_transfer_config: ECTransferConfig | None = None
108
+ """The configurations for distributed EC cache transfer."""
109
# some opaque config, only used to provide additional information
110
# for the hash computation, mainly used for testing, debugging or out of
111
# tree config registration.
@@ -183,6 +186,10 @@ class VllmConfig:
186
vllm_factors.append(self.kv_transfer_config.compute_hash())
187
else:
188
vllm_factors.append("None")
189
+ if self.ec_transfer_config:
190
+ vllm_factors.append(self.ec_transfer_config.compute_hash())
191
+ else:
192
+ vllm_factors.append("None")
193
if self.additional_config:
194
if isinstance(additional_config := self.additional_config, dict):
195
additional_config_hash = hashlib.md5(
vllm/distributed/ec_transfer/__init__.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.distributed.ec_transfer.ec_transfer_state import (
5
+ ensure_ec_transfer_initialized,
6
+ get_ec_transfer,
7
+ has_ec_transfer,
8
+)
9
+
10
+__all__ = [
11
+ "get_ec_transfer",
12
+ "ensure_ec_transfer_initialized",
13
+ "has_ec_transfer",
14
+]
vllm/distributed/ec_transfer/ec_connector/__init__.py
vllm/distributed/ec_transfer/ec_connector/base.py
new
+247
@@ -0,0 +1,247 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+ECConnectorBase Class for Distributed Encoder Cache &
5
+P2P Encoder cache communication in V1
6
+
7
+The class provides the following primitives:
8
+ Scheduler-side: runs in the scheduler, binds metadata, which
9
+ is used by the worker-side to load/save Encoder cache.
10
+ check_caches_exist() - Check whether Encoder cache of requests exist
11
+ update_state_after_alloc() - update ECConnector state after
12
+ allocate. This will decide to load the cache or not
13
+ request_finished() - called when a request is finished,
14
+ free the cache with the requests
15
+
16
+ Worker-side: runs in each worker, loads/saves Encoder Cache to/from
17
+ the Connector based on the metadata.
18
+ start_load_ec() - starts loading all ECs (maybe async)
19
+ wait_for_save() - blocks until all saves are done
20
+
21
+ get_finished() - called with ids of finished requests, returns
22
+ ids of requests that have completed async sending/recving.
23
+"""
24
+
25
+import enum
26
+from abc import ABC, abstractmethod
27
+from typing import TYPE_CHECKING, Any
28
+
29
+import torch
30
+
31
+from vllm.logger import init_logger
32
+from vllm.v1.core.sched.output import SchedulerOutput
33
+from vllm.v1.outputs import ECConnectorOutput
34
+
35
+if TYPE_CHECKING:
36
+ from vllm.config import VllmConfig
37
+ from vllm.v1.request import Request
38
+
39
+logger = init_logger(__name__)
40
+
41
+
42
+class ECConnectorRole(enum.Enum):
43
+ # Connector running in the scheduler process
44
+ SCHEDULER = 0
45
+
46
+ # Connector running in the worker process
47
+ WORKER = 1
48
+
49
+
50
+class ECConnectorMetadata(ABC): # noqa: B024
51
+ """
52
+ Abstract Metadata used to communicate between the
53
+ Scheduler ECConnector and Worker ECConnector.
54
+ """
55
+
56
+ pass
57
+
58
+
59
+class ECConnectorBase(ABC):
60
+ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
61
+ self._connector_metadata: ECConnectorMetadata | None = None
62
+ self._vllm_config = vllm_config
63
+ self._role = role
64
+ if vllm_config.ec_transfer_config is not None:
65
+ self._is_producer = vllm_config.ec_transfer_config.is_ec_producer
66
+ else:
67
+ raise ValueError("ec_transfer_config must be set for ECConnectorBase")
68
+
69
+ @property
70
+ def role(self) -> ECConnectorRole:
71
+ return self._role
72
+
73
+ @property
74
+ def is_producer(self) -> bool:
75
+ return self._is_producer
76
+
77
+ # ==============================
78
+ # Worker-side methods
79
+ # ==============================
80
+
81
+ def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> None:
82
+ """Set the connector metadata from the scheduler.
83
+
84
+ This function should be called by the model runner every time
85
+ before the model execution. The metadata will be used for runtime
86
+ EC cache loading.
87
+
88
+ Args:
89
+ connector_metadata (dict): the connector metadata.
90
+ """
91
+ self._connector_metadata = connector_metadata
92
+
93
+ def clear_connector_metadata(self) -> None:
94
+ """Clear the connector metadata.
95
+
96
+ This function should be called by the model runner every time
97
+ after the model execution.
98
+ """
99
+ self._connector_metadata = None
100
+
101
+ def _get_connector_metadata(self) -> ECConnectorMetadata:
102
+ """Get the connector metadata.
103
+
104
+ This function should only be called inside the connector.
105
+
106
+ Returns:
107
+ ConnectorMetadata: the connector metadata.
108
+ """
109
+
110
+ # Should only be called while set to valid metadata.
111
+ assert self._connector_metadata is not None
112
+ return self._connector_metadata
113
+
114
+ def register_caches(
115
+ self,
116
+ ec_caches: dict[str, torch.Tensor],
117
+ ):
118
+ """
119
+ Initialize with the EC caches.
120
+ Args:
121
+ ec_caches: dictionary of encoder cache
122
+ """
123
+ # TODO: Implement this later for P2P feature
124
+ return
125
+
126
+ @abstractmethod
127
+ def start_load_caches(
128
+ self, encoder_cache: dict[str, torch.Tensor], **kwargs
129
+ ) -> None:
130
+ """
131
+ Start loading the cache from the connector into vLLM's encoder cache.
132
+
133
+ This method loads the encoder cache based on metadata provided by the scheduler.
134
+ It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
135
+ the `encoder_cache` and `mm_hash` are stored in `kwargs`.
136
+
137
+ Args:
138
+ encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
139
+ data hashes (`mm_hash`) to encoder cache tensors.
140
+ kwargs (dict): Additional keyword arguments for the connector.
141
+ """
142
+ pass
143
+
144
+ @abstractmethod
145
+ def save_caches(
146
+ self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs
147
+ ) -> None:
148
+ """
149
+ Save the encoder cache to the connector.
150
+
151
+ This method saves the encoder cache from the worker's local storage
152
+ to shared storage or another external connector.
153
+
154
+ Args:
155
+ encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
156
+ data hashes (`mm_hash`) to encoder cache tensors.
157
+ mm_hash (str): The hash of the multimodal data whose cache is being saved.
158
+ kwargs (dict): Additional keyword arguments for the connector.
159
+ """
160
+ pass
161
+
162
+ def get_finished(
163
+ self, finished_req_ids: set[str]
164
+ ) -> tuple[set[str] | None, set[str] | None]:
165
+ """
166
+ Notifies worker-side connector ids of requests that have
167
+ finished generating tokens on the worker.
168
+ The scheduler process (via the Executors) will use this output
169
+ to track which workers are done.
170
+
171
+ Returns:
172
+ ids of requests that have finished asynchronous transfer
173
+ (requests that previously returned True from request_finished()),
174
+ tuple of (sending/saving ids, recving/loading ids).
175
+ The finished saves/sends req ids must belong to a set provided in a
176
+ call to this method (this call or a prior one).
177
+ """
178
+ return None, None
179
+
180
+ # ==============================
181
+ # Scheduler-side methods
182
+ # ==============================
183
+
184
+ @abstractmethod
185
+ def has_caches(
186
+ self,
187
+ request: "Request",
188
+ ) -> list[bool]:
189
+ """
190
+ Check if encoder cache exists for each mm data of requests
191
+
192
+ Args:
193
+ request (Request): the request object.
194
+
195
+ Returns:
196
+ A list bool where ith value is True if cache exist for
197
+ ith mm_data of requests
198
+ """
199
+ pass
200
+
201
+ @abstractmethod
202
+ def update_state_after_alloc(self, request: "Request", index: int):
203
+ """
204
+ Update ECConnector state to decide allocate cache for requests
205
+
206
+ Args:
207
+ request (Request): the request object.
208
+ """
209
+ pass
210
+
211
+ @abstractmethod
212
+ def build_connector_meta(
213
+ self, scheduler_output: SchedulerOutput
214
+ ) -> ECConnectorMetadata:
215
+ """
216
+ Build the connector metadata for this step.
217
+
218
+ This function should NOT modify fields in the scheduler_output.
219
+ Also, calling this function will reset the state of the connector.
220
+
221
+ Args:
222
+ scheduler_output (SchedulerOutput): the scheduler output object.
223
+ """
224
+ pass
225
+
226
+ def update_connector_output(self, connector_output: ECConnectorOutput):
227
+ """
228
+ Update ECConnector state from worker-side connectors output.
229
+
230
+ Args:
231
+ connector_output (ECConnectorOutput): the worker-side
232
+ connectors output.
233
+ """
234
+ return
235
+
236
+ def request_finished(
237
+ self, request: "Request"
238
+ ) -> tuple[bool, dict[str, Any] | None]:
239
+ """
240
+ Called when a request has finished, before its encoder cache is freed.
241
+
242
+ Returns:
243
+ True if the request is being saved/sent asynchronously and cached
244
+ should not be freed until the request_id is returned from
245
+ get_finished().
246
+ """
247
+ return False, None
vllm/distributed/ec_transfer/ec_connector/factory.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib
5
+from collections.abc import Callable
6
+from typing import TYPE_CHECKING
7
+
8
+# yapf: disable
9
+from vllm.distributed.ec_transfer.ec_connector.base import (
10
+ ECConnectorBase,
11
+ ECConnectorRole,
12
+)
13
+from vllm.logger import init_logger
14
+
15
+# yapf: enable
16
+
17
+if TYPE_CHECKING:
18
+ from vllm.config import ECTransferConfig, VllmConfig
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+class ECConnectorFactory:
24
+ _registry: dict[str, Callable[[], type[ECConnectorBase]]] = {}
25
+
26
+ @classmethod
27
+ def register_connector(cls, name: str, module_path: str, class_name: str) -> None:
28
+ """Register a connector with a lazy-loading module and class name."""
29
+ if name in cls._registry:
30
+ raise ValueError(f"Connector '{name}' is already registered.")
31
+
32
+ def loader() -> type[ECConnectorBase]:
33
+ module = importlib.import_module(module_path)
34
+ return getattr(module, class_name)
35
+
36
+ cls._registry[name] = loader
37
+
38
+ @classmethod
39
+ def create_connector(
40
+ cls,
41
+ config: "VllmConfig",
42
+ role: ECConnectorRole,
43
+ ) -> ECConnectorBase:
44
+ ec_transfer_config = config.ec_transfer_config
45
+ if ec_transfer_config is None:
46
+ raise ValueError("ec_transfer_config must be set to create a connector")
47
+ connector_cls = cls.get_connector_class(ec_transfer_config)
48
+ logger.info(
49
+ "Creating connector with name: %s and engine_id: %s",
50
+ connector_cls.__name__,
51
+ ec_transfer_config.engine_id,
52
+ )
53
+ # Connector is explicitly separated into two roles.
54
+ # Scheduler connector:
55
+ # - Co-locate with scheduler process
56
+ # - Should only be used inside the Scheduler class
57
+ # Worker connector:
58
+ # - Co-locate with worker process
59
+ return connector_cls(config, role)
60
+
61
+ @classmethod
62
+ def get_connector_class(
63
+ cls, ec_transfer_config: "ECTransferConfig"
64
+ ) -> type[ECConnectorBase]:
65
+ """Get the connector class by name."""
66
+ connector_name = ec_transfer_config.ec_connector
67
+ if connector_name is None:
68
+ raise ValueError("EC connect must not be None")
69
+ elif connector_name in cls._registry:
70
+ connector_cls = cls._registry[connector_name]()
71
+ else:
72
+ connector_module_path = ec_transfer_config.ec_connector_module_path
73
+ if connector_module_path is None:
74
+ raise ValueError(f"Unsupported connector type: {connector_name}")
75
+ connector_module = importlib.import_module(connector_module_path)
76
+ connector_cls = getattr(connector_module, connector_name)
77
+ return connector_cls
78
+
79
+
80
+# Register various connectors here.
81
+# The registration should not be done in each individual file, as we want to
82
+# only load the files corresponding to the current connector.
83
+
84
+ECConnectorFactory.register_connector(
85
+ "ECSharedStorageConnector",
86
+ "vllm.distributed.ec_transfer.ec_connector.shared_storage_connector",
87
+ "ECSharedStorageConnector",
88
+)
vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py
new
+201
@@ -0,0 +1,201 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+from dataclasses import dataclass
5
+from typing import TYPE_CHECKING
6
+
7
+import safetensors
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.distributed.ec_transfer.ec_connector.base import (
11
+ ECConnectorBase,
12
+ ECConnectorMetadata,
13
+ ECConnectorRole,
14
+)
15
+from vllm.logger import init_logger
16
+from vllm.v1.core.sched.output import SchedulerOutput
17
+
18
+if TYPE_CHECKING:
19
+ from vllm.v1.request import Request
20
+
21
+logger = init_logger(__name__)
22
+
23
+
24
+@dataclass
25
+class MMMeta:
26
+ mm_hash: str
27
+ num_token: int
28
+
29
+ @staticmethod
30
+ def make_meta(mm_hash, num_token) -> "MMMeta":
31
+ return MMMeta(mm_hash=mm_hash, num_token=num_token)
32
+
33
+
34
+@dataclass
35
+class ECSharedStorageConnectorMetadata(ECConnectorMetadata):
36
+ mm_datas: list[MMMeta]
37
+
38
+ def __init__(self):
39
+ self.mm_datas = []
40
+
41
+ def add_mm_data(self, mm_data: MMMeta):
42
+ self.mm_datas.append(mm_data)
43
+
44
+
45
+class ECSharedStorageConnector(ECConnectorBase):
46
+ # NOTE: This is Simple debug implementation of the EC connector.
47
+ # It save / load the EC cache to / from the disk.
48
+
49
+ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
50
+ super().__init__(vllm_config=vllm_config, role=role)
51
+ # req_id -> index
52
+ self._mm_datas_need_loads: dict[str, int] = {}
53
+ transfer_config = vllm_config.ec_transfer_config
54
+ if transfer_config is not None:
55
+ self._storage_path = transfer_config.get_from_extra_config(
56
+ "shared_storage_path", "/tmp"
57
+ )
58
+ logger.debug(transfer_config)
59
+ logger.debug("Shared storage path is %s", self._storage_path)
60
+ else:
61
+ raise ValueError("ec_transfer_config must be set for ECConnectorBase")
62
+
63
+ def start_load_caches(self, encoder_cache, **kwargs) -> None:
64
+ """
65
+ Start loading the cache from the connector into vLLM's encoder cache.
66
+
67
+ This method loads the encoder cache based on metadata provided by the scheduler.
68
+ It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
69
+ the `encoder_cache` and `mm_hash` are stored in `kwargs`.
70
+
71
+ Args:
72
+ encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
73
+ data hashes (`mm_hash`) to encoder cache tensors.
74
+ kwargs (dict): Additional keyword arguments for the connector.
75
+ """
76
+
77
+ # Get the metadata
78
+ metadata: ECConnectorMetadata = self._get_connector_metadata()
79
+ assert isinstance(metadata, ECSharedStorageConnectorMetadata)
80
+ assert encoder_cache is not None
81
+ if metadata is None:
82
+ logger.warning(
83
+ (
84
+ "In connector.start_load_caches, ",
85
+ "but the connector metadata is None",
86
+ )
87
+ )
88
+ return
89
+ # Load the EC for each mm data
90
+ for mm_data in metadata.mm_datas:
91
+ if mm_data.mm_hash in encoder_cache:
92
+ continue
93
+ filename = self._generate_filename_debug(mm_data.mm_hash)
94
+ ec_cache = safetensors.torch.load_file(filename)["ec_cache"].cuda()
95
+ encoder_cache[mm_data.mm_hash] = ec_cache
96
+ logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
97
+
98
+ def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None:
99
+ """
100
+ Save the encoder cache to the connector.
101
+
102
+ This method saves the encoder cache from the worker's local storage
103
+ to shared storage or another external connector.
104
+
105
+ Args:
106
+ encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
107
+ data hashes (`mm_hash`) to encoder cache tensors.
108
+ mm_hash (str): The hash of the multimodal data whose cache is being saved.
109
+ kwargs (dict): Additional keyword arguments for the connector.
110
+ """
111
+ # Return if it is PD Instance
112
+ if not self.is_producer:
113
+ return
114
+ filename = self._generate_filename_debug(mm_hash)
115
+ ec_cache = encoder_cache[mm_hash]
116
+ tensors = {"ec_cache": ec_cache.detach().cpu()}
117
+ safetensors.torch.save_file(tensors, filename)
118
+ logger.debug("Save cache successful for mm_hash %s", mm_hash)
119
+
120
+ def has_caches(
121
+ self,
122
+ request: "Request",
123
+ ) -> list[bool]:
124
+ """
125
+ Check if cache exist externally for each mm_data of request
126
+
127
+ Args:
128
+ request (Request): the request object.
129
+
130
+ Returns:
131
+ List of bool indicate that ith mm_data exist in cache or not
132
+ """
133
+ result = []
134
+ for feature in request.mm_features:
135
+ result.append(self._found_match_for_mm_data(feature.identifier))
136
+ return result
137
+
138
+ def update_state_after_alloc(
139
+ self,
140
+ request: "Request",
141
+ index: int,
142
+ ) -> None:
143
+ """
144
+ Update ECConnector state after encoder cache allocation.
145
+ """
146
+ mm_hash = request.mm_features[index].identifier
147
+ num_encoder_token = request.get_num_encoder_tokens(index)
148
+ # Insert mm_hash only if this block has not been recorded yet.
149
+ self._mm_datas_need_loads[mm_hash] = num_encoder_token
150
+
151
+ def build_connector_meta(
152
+ self,
153
+ scheduler_output: SchedulerOutput,
154
+ ) -> ECConnectorMetadata:
155
+ """Build the connector metadata for this step.
156
+
157
+ This function should NOT modify any fields in the scheduler_output.
158
+ Also, calling this function will reset the state of the connector.
159
+ This only build for load mm_data only
160
+ Args:
161
+ scheduler_output (SchedulerOutput): the scheduler output object.
162
+ """
163
+ meta = ECSharedStorageConnectorMetadata()
164
+ for mm_hash, num_encoder_token in self._mm_datas_need_loads.items():
165
+ meta.add_mm_data(MMMeta.make_meta(mm_hash, num_encoder_token))
166
+ self._mm_datas_need_loads.clear()
167
+ return meta
168
+
169
+ # ==============================
170
+ # Helper functions
171
+ # ==============================
172
+
173
+ def _found_match_for_mm_data(self, mm_hash) -> bool:
174
+ """Check if the cache is hit for the request."""
175
+ filename = self._generate_filename_debug(mm_hash)
176
+ return os.path.exists(filename)
177
+
178
+ def _generate_foldername_debug(
179
+ self,
180
+ mm_hash: str,
181
+ create_folder: bool = True, # <- now defaults to True
182
+ ) -> str:
183
+ """
184
+ Return the folder in which the cache for this mm_hash lives.
185
+ If `create_folder` is True (default) the directory is created
186
+ recursively the first time it is needed.
187
+ """
188
+ foldername = os.path.join(self._storage_path, mm_hash)
189
+ if create_folder:
190
+ os.makedirs(foldername, exist_ok=True)
191
+ return foldername
192
+
193
+ def _generate_filename_debug(self, mm_hash: str) -> str:
194
+ """
195
+ Return the full path of the safetensors file for this mm_hash.
196
+ Ensures the parent directory exists because
197
+ `_generate_foldername_debug` is called with its default
198
+ (`create_folder=True`).
199
+ """
200
+ foldername = self._generate_foldername_debug(mm_hash) # <- folder auto-created
201
+ return os.path.join(foldername, "encoder_cache.safetensors")
vllm/distributed/ec_transfer/ec_transfer_state.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import TYPE_CHECKING
4
+
5
+from vllm import envs
6
+from vllm.distributed.ec_transfer.ec_connector.base import (
7
+ ECConnectorBase,
8
+ ECConnectorRole,
9
+)
10
+from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
11
+
12
+if TYPE_CHECKING:
13
+ from vllm.config import VllmConfig
14
+
15
+_EC_CONNECTOR_AGENT: ECConnectorBase | None = None
16
+
17
+
18
+def get_ec_transfer() -> ECConnectorBase:
19
+ assert _EC_CONNECTOR_AGENT is not None, "disaggregated EC cache is not initialized"
20
+ return _EC_CONNECTOR_AGENT
21
+
22
+
23
+def has_ec_transfer() -> bool:
24
+ return _EC_CONNECTOR_AGENT is not None
25
+
26
+
27
+def ensure_ec_transfer_initialized(vllm_config: "VllmConfig") -> None:
28
+ """
29
+ Initialize EC cache connector.
30
+ """
31
+
32
+ global _EC_CONNECTOR_AGENT
33
+
34
+ if vllm_config.ec_transfer_config is None:
35
+ return
36
+
37
+ if (
38
+ vllm_config.ec_transfer_config.is_ec_transfer_instance
39
+ and _EC_CONNECTOR_AGENT is None
40
+ ):
41
+ if envs.VLLM_USE_V1:
42
+ _EC_CONNECTOR_AGENT = ECConnectorFactory.create_connector(
43
+ config=vllm_config, role=ECConnectorRole.WORKER
44
+ )
45
+ else:
46
+ raise ValueError("V0 is no longer supported")
vllm/engine/arg_utils.py
+7
@@ -38,6 +38,7 @@ from vllm.config import (
38
CompilationConfig,
39
ConfigType,
40
DeviceConfig,
41
+ ECTransferConfig,
42
EPLBConfig,
43
KVEventsConfig,
44
KVTransferConfig,
@@ -527,6 +528,8 @@ class EngineArgs:
528
kv_transfer_config: KVTransferConfig | None = None
529
kv_events_config: KVEventsConfig | None = None
530
531
+ ec_transfer_config: ECTransferConfig | None = None
532
+
533
generation_config: str = ModelConfig.generation_config
534
enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
535
override_generation_config: dict[str, Any] = get_field(
@@ -1105,6 +1108,9 @@ class EngineArgs:
1108
"--kv-transfer-config", **vllm_kwargs["kv_transfer_config"]
1109
)
1110
vllm_group.add_argument("--kv-events-config", **vllm_kwargs["kv_events_config"])
1111
+ vllm_group.add_argument(
1112
+ "--ec-transfer-config", **vllm_kwargs["ec_transfer_config"]
1113
+ )
1114
vllm_group.add_argument(
1115
"--compilation-config", "-O", **vllm_kwargs["compilation_config"]
1116
)
@@ -1676,6 +1682,7 @@ class EngineArgs:
1682
compilation_config=self.compilation_config,
1683
kv_transfer_config=self.kv_transfer_config,
1684
kv_events_config=self.kv_events_config,
1685
+ ec_transfer_config=self.ec_transfer_config,
1686
additional_config=self.additional_config,
1687
)
1688
vllm/model_executor/warmup/kernel_warmup.py
+12
-4
@@ -49,10 +49,18 @@ def kernel_warmup(worker: "Worker"):
49
except NotImplementedError:
50
return False
51
52
- if not worker.model_runner.is_pooling_model and all(
53
- _is_flashinfer_backend(group.backend)
54
- for groups in worker.model_runner.attn_groups
55
- for group in groups
52
+ # NOTE: we add check for empty attn_groups to avoid errors when
53
+ # deploying models such as E instances and encoder-only models.
54
+ # As for those models, worker.model_runner.attn_groups is empty.
55
+ # This change is made during EPD feature development.
56
+ if (
57
+ not worker.model_runner.is_pooling_model
58
+ and worker.model_runner.attn_groups
59
+ and all(
60
+ _is_flashinfer_backend(group.backend)
61
+ for groups in worker.model_runner.attn_groups
62
+ for group in groups
63
+ )
64
):
65
logger.info("Warming up FlashInfer attention.")
66
# Warmup with mixed batch containing both prefill and decode tokens
vllm/v1/core/sched/output.py
+5
@@ -14,6 +14,7 @@ if TYPE_CHECKING:
14
import numpy.typing as npt
15
import torch
16
17
+ from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
18
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
19
from vllm.lora.request import LoRARequest
20
from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -21,6 +22,7 @@ if TYPE_CHECKING:
22
from vllm.sampling_params import SamplingParams
23
from vllm.v1.request import Request
24
else:
25
+ ECConnectorMetadata = object
26
KVConnectorMetadata = object
27
LoRARequest = object
28
MultiModalFeatureSpec = object
@@ -188,6 +190,9 @@ class SchedulerOutput:
190
# KV Cache Connector metadata.
191
kv_connector_metadata: KVConnectorMetadata | None = None
192
193
+ # EC Cache Connector metadata
194
+ ec_connector_metadata: ECConnectorMetadata | None = None
195
+
196
197
@dataclass
198
class GrammarOutput:
vllm/v1/core/sched/scheduler.py
+52
-4
@@ -7,6 +7,11 @@ from collections.abc import Iterable
7
from typing import Any
8
9
from vllm.config import VllmConfig
10
+from vllm.distributed.ec_transfer.ec_connector.base import (
11
+ ECConnectorMetadata,
12
+ ECConnectorRole,
13
+)
14
+from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
15
from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
16
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
17
from vllm.distributed.kv_transfer.kv_connector.v1 import (
@@ -14,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
19
KVConnectorRole,
20
SupportsHMA,
21
)
22
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
23
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
24
from vllm.logger import init_logger
25
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
@@ -104,6 +110,11 @@ class Scheduler(SchedulerInterface):
110
self.kv_events_config,
111
self.parallel_config.data_parallel_rank,
112
)
113
+ self.ec_connector = None
114
+ if self.vllm_config.ec_transfer_config is not None:
115
+ self.ec_connector = ECConnectorFactory.create_connector(
116
+ config=self.vllm_config, role=ECConnectorRole.SCHEDULER
117
+ )
118
119
num_gpu_blocks = self.cache_config.num_gpu_blocks
120
assert num_gpu_blocks is not None and num_gpu_blocks > 0
@@ -230,12 +241,14 @@ class Scheduler(SchedulerInterface):
241
242
# Schedule encoder inputs.
243
encoder_inputs_to_schedule = None
244
+ external_load_encoder_input: list[int] = []
245
new_encoder_compute_budget = encoder_compute_budget
246
if request.has_encoder_inputs:
247
(
248
encoder_inputs_to_schedule,
249
num_new_tokens,
250
new_encoder_compute_budget,
251
+ external_load_encoder_input,
252
) = self._try_schedule_encoder_inputs(
253
request,
254
request.num_computed_tokens,
@@ -342,6 +355,11 @@ class Scheduler(SchedulerInterface):
355
for i in encoder_inputs_to_schedule:
356
self.encoder_cache_manager.allocate(request, i)
357
encoder_compute_budget = new_encoder_compute_budget
358
+ if external_load_encoder_input:
359
+ for i in external_load_encoder_input:
360
+ self.encoder_cache_manager.allocate(request, i)
361
+ if self.ec_connector is not None:
362
+ self.ec_connector.update_state_after_alloc(request, i)
363
364
# Record the LoRAs in scheduled_running_reqs
365
scheduled_loras: set[int] = set()
@@ -445,6 +463,7 @@ class Scheduler(SchedulerInterface):
463
num_computed_tokens = request.num_computed_tokens
464
465
encoder_inputs_to_schedule = None
466
+ external_load_encoder_input = []
467
new_encoder_compute_budget = encoder_compute_budget
468
469
# KVTransfer: loading remote KV, do not allocate for new work.
@@ -480,6 +499,7 @@ class Scheduler(SchedulerInterface):
499
encoder_inputs_to_schedule,
500
num_new_tokens,
501
new_encoder_compute_budget,
502
+ external_load_encoder_input,
503
) = self._try_schedule_encoder_inputs(
504
request,
505
num_computed_tokens,
@@ -583,7 +603,12 @@ class Scheduler(SchedulerInterface):
603
for i in encoder_inputs_to_schedule:
604
self.encoder_cache_manager.allocate(request, i)
605
encoder_compute_budget = new_encoder_compute_budget
586
-
606
+ # Allocate for external load encoder cache
607
+ if external_load_encoder_input:
608
+ for i in external_load_encoder_input:
609
+ self.encoder_cache_manager.allocate(request, i)
610
+ if self.ec_connector is not None:
611
+ self.ec_connector.update_state_after_alloc(request, i)
612
# Put back any skipped requests at the head of the waiting queue
613
if skipped_waiting_requests:
614
self.waiting.prepend_requests(skipped_waiting_requests)
@@ -591,6 +616,7 @@ class Scheduler(SchedulerInterface):
616
# Check if the scheduling constraints are satisfied.
617
total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
618
assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
619
+
620
assert token_budget >= 0
621
assert len(self.running) <= self.max_num_running_reqs
622
# Since some requests in the RUNNING queue may not be scheduled in
@@ -653,8 +679,18 @@ class Scheduler(SchedulerInterface):
679
# 2. Wrap up all the KV cache load / save ops into an opaque object
680
# 3. Clear the internal states of the connector
681
if self.connector is not None:
656
- meta = self.connector.build_connector_meta(scheduler_output)
682
+ meta: KVConnectorMetadata = self.connector.build_connector_meta(
683
+ scheduler_output
684
+ )
685
scheduler_output.kv_connector_metadata = meta
686
+
687
+ # Build the connector meta for ECConnector
688
+ if self.ec_connector is not None:
689
+ ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta(
690
+ scheduler_output
691
+ )
692
+ scheduler_output.ec_connector_metadata = ec_meta
693
+
694
with record_function_or_nullcontext("schedule: update_after_schedule"):
695
self._update_after_schedule(scheduler_output)
696
return scheduler_output
@@ -755,7 +791,7 @@ class Scheduler(SchedulerInterface):
791
num_computed_tokens: int,
792
num_new_tokens: int,
793
encoder_compute_budget: int,
758
- ) -> tuple[list[int], int, int]:
794
+ ) -> tuple[list[int], int, int, list[int]]:
795
"""
796
Determine which encoder inputs need to be scheduled in the current step,
797
and update `num_new_tokens` and encoder token budget accordingly.
@@ -765,6 +801,7 @@ class Scheduler(SchedulerInterface):
801
in this step, i.e.,
802
[num_computed_tokens, num_computed_tokens + num_new_tokens).
803
- It is not already computed and stored in the encoder cache.
804
+ - It is not exist on remote encoder cache (via ECConnector)
805
- There is sufficient encoder token budget to process it.
806
- The encoder cache has space to store it.
807
@@ -776,12 +813,16 @@ class Scheduler(SchedulerInterface):
813
blocks and externally cached blocks (via KVConnector).
814
"""
815
if num_new_tokens == 0 or not request.has_encoder_inputs:
779
- return [], num_new_tokens, encoder_compute_budget
816
+ return [], num_new_tokens, encoder_compute_budget, []
817
encoder_inputs_to_schedule: list[int] = []
818
mm_features = request.mm_features
819
assert mm_features is not None
820
assert len(mm_features) > 0
821
+ external_load_encoder_input = []
822
823
+ # Check remote cache first
824
+ if self.ec_connector is not None:
825
+ remote_cache_has_item = self.ec_connector.has_caches(request)
826
# NOTE: since scheduler operates on the request level (possibly with
827
# multiple encoder inputs per request), we need to create temporary
828
# trackers for accounting at the encoder input level.
@@ -862,6 +903,12 @@ class Scheduler(SchedulerInterface):
903
num_new_tokens = 0
904
break
905
906
+ if self.ec_connector is not None and remote_cache_has_item[i]:
907
+ mm_hashes_to_schedule.add(request.mm_features[i].identifier)
908
+ external_load_encoder_input.append(i)
909
+ num_tokens_to_schedule += num_encoder_tokens
910
+ continue
911
+
912
num_tokens_to_schedule += num_encoder_tokens
913
encoder_compute_budget -= num_encoder_tokens
914
mm_hashes_to_schedule.add(request.mm_features[i].identifier)
@@ -871,6 +918,7 @@ class Scheduler(SchedulerInterface):
918
encoder_inputs_to_schedule,
919
num_new_tokens,
920
encoder_compute_budget,
921
+ external_load_encoder_input,
922
)
923
924
def get_grammar_bitmask(
vllm/v1/outputs.py
+46
@@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, NamedTuple
8
import numpy as np
9
import torch
10
11
+from vllm.v1.core.sched.output import SchedulerOutput
12
+
13
if TYPE_CHECKING:
14
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
15
else:
@@ -136,6 +138,13 @@ class KVConnectorOutput:
138
)
139
140
141
+@dataclass
142
+class ECConnectorOutput:
143
+ # [mm_hash]
144
+ finished_sending: set[str] | None = None
145
+ finished_recving: set[str] | None = None
146
+
147
+
148
# ModelRunnerOutput is serialized and sent to the scheduler process.
149
# This is expensive for torch.Tensor so prefer to use list instead.
150
@dataclass
@@ -167,6 +176,8 @@ class ModelRunnerOutput:
176
177
kv_connector_output: KVConnectorOutput | None = None
178
179
+ ec_connector_output: ECConnectorOutput | None = None
180
+
181
# req_id -> num_nans_in_logits
182
num_nans_in_logits: dict[str, int] | None = None
183
@@ -192,6 +203,41 @@ class DraftTokenIds:
203
draft_token_ids: list[list[int]]
204
205
206
+def make_empty_encoder_model_runner_output(
207
+ scheduler_output: "SchedulerOutput",
208
+) -> ModelRunnerOutput:
209
+ """
210
+ Create a ModelRunnerOutput stub that contains the correct
211
+ per-request bookkeeping but no generated data yet.
212
+ """
213
+ if not scheduler_output.num_scheduled_tokens:
214
+ return EMPTY_MODEL_RUNNER_OUTPUT
215
+
216
+ # Convert to list so we get a deterministic, indexable sequence
217
+ req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys())
218
+
219
+ # Give every request its own contiguous index
220
+ req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)}
221
+
222
+ # No tokens generated yet ⇒ one empty list per request
223
+ sampled_token_ids: list[list[int]] = [[0] for _ in req_ids]
224
+
225
+ # Pooler outputs are not available yet ⇒ use None placeholders
226
+ pooler_output: list[torch.Tensor | None] = [None for _ in req_ids]
227
+
228
+ return ModelRunnerOutput(
229
+ req_ids=req_ids,
230
+ req_id_to_index=req_id_to_index,
231
+ sampled_token_ids=sampled_token_ids,
232
+ logprobs=None,
233
+ prompt_logprobs_dict={},
234
+ pooler_output=pooler_output,
235
+ kv_connector_output=None,
236
+ ec_connector_output=None,
237
+ num_nans_in_logits=None,
238
+ )
239
+
240
+
241
EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
242
req_ids=[],
243
req_id_to_index={},
vllm/v1/worker/ec_connector_model_runner_mixin.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Define EC connector functionality mixin for model runners.
5
+"""
6
+
7
+from collections.abc import Generator
8
+from contextlib import AbstractContextManager, contextmanager, nullcontext
9
+from typing import (
10
+ TYPE_CHECKING, # noqa: UP035
11
+)
12
+
13
+import torch
14
+
15
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
16
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase
17
+from vllm.logger import init_logger
18
+from vllm.v1.outputs import ECConnectorOutput
19
+
20
+if TYPE_CHECKING:
21
+ from vllm.v1.core.sched.output import SchedulerOutput
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+# Defined as a EC connector functionality mixin for ModelRunner (GPU, TPU)
27
+class ECConnectorModelRunnerMixin:
28
+ @staticmethod
29
+ def maybe_save_ec_to_connector(
30
+ encoder_cache: dict[str, torch.Tensor],
31
+ mm_hash: str,
32
+ ):
33
+ if not has_ec_transfer():
34
+ logger.debug("Not have ec transfer please check")
35
+ return
36
+ connector = get_ec_transfer()
37
+ connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash)
38
+
39
+ @staticmethod
40
+ def get_finished_ec_transfers(
41
+ scheduler_output: "SchedulerOutput",
42
+ ) -> tuple[set[str] | None, set[str] | None]:
43
+ if has_ec_transfer():
44
+ return get_ec_transfer().get_finished(scheduler_output.finished_req_ids)
45
+ return None, None
46
+
47
+ @staticmethod
48
+ def maybe_get_ec_connector_output(
49
+ scheduler_output: "SchedulerOutput",
50
+ encoder_cache: dict[str, torch.Tensor],
51
+ **kwargs,
52
+ ) -> AbstractContextManager[ECConnectorOutput | None]:
53
+ return (
54
+ ECConnectorModelRunnerMixin._get_ec_connector_output(
55
+ scheduler_output, encoder_cache, **kwargs
56
+ )
57
+ if has_ec_transfer()
58
+ else nullcontext()
59
+ )
60
+
61
+ # This context manager must be used within an active forward context.
62
+ # It encapsulates the entire EC conector lifecycle within execute_model
63
+ @staticmethod
64
+ @contextmanager
65
+ def _get_ec_connector_output(
66
+ scheduler_output: "SchedulerOutput",
67
+ encoder_cache: dict[str, torch.Tensor],
68
+ **kwargs,
69
+ ) -> Generator[ECConnectorOutput, None, None]:
70
+ output = ECConnectorOutput()
71
+
72
+ ec_connector = get_ec_transfer()
73
+ assert isinstance(ec_connector, ECConnectorBase)
74
+ assert scheduler_output.ec_connector_metadata is not None
75
+ ec_connector.bind_connector_metadata(scheduler_output.ec_connector_metadata)
76
+
77
+ if not ec_connector.is_producer:
78
+ ec_connector.start_load_caches(encoder_cache, **kwargs)
79
+
80
+ try:
81
+ yield output
82
+ finally:
83
+ output.finished_sending, output.finished_recving = (
84
+ ec_connector.get_finished(scheduler_output.finished_req_ids)
85
+ )
86
+
87
+ ec_connector.clear_connector_metadata()
vllm/v1/worker/gpu_model_runner.py
+36
-4
@@ -35,6 +35,7 @@ from vllm.config import (
35
get_layers_from_vllm_config,
36
update_config,
37
)
38
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
39
from vllm.distributed.eplb.eplb_state import EplbState
40
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
41
from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
@@ -114,12 +115,14 @@ from vllm.v1.outputs import (
115
EMPTY_MODEL_RUNNER_OUTPUT,
116
AsyncModelRunnerOutput,
117
DraftTokenIds,
118
+ ECConnectorOutput,
119
KVConnectorOutput,
120
LogprobsLists,
121
LogprobsTensors,
122
ModelRunnerOutput,
123
PoolerOutput,
124
SamplerOutput,
125
+ make_empty_encoder_model_runner_output,
126
)
127
from vllm.v1.pool.metadata import PoolingMetadata
128
from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs
@@ -134,6 +137,7 @@ from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
137
from vllm.v1.structured_output.utils import apply_grammar_bitmask
138
from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
139
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
140
+from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
141
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
142
from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
143
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -237,9 +241,12 @@ class ExecuteModelState(NamedTuple):
241
sample_hidden_states: torch.Tensor
242
aux_hidden_states: list[torch.Tensor] | None
243
kv_connector_output: KVConnectorOutput | None
244
+ ec_connector_output: ECConnectorOutput | None
245
246
242
-class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
247
+class GPUModelRunner(
248
+ LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin
249
+):
250
def __init__(
251
self,
252
vllm_config: VllmConfig,
@@ -1873,6 +1880,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1880
output,
1881
is_embed=pos_info.is_embed,
1882
)
1883
+ logger.debug("Finish execute for mm hash %s", mm_hash)
1884
+ self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)
1885
1886
def _gather_mm_embeddings(
1887
self,
@@ -2191,20 +2200,27 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2200
torch.Tensor,
2201
IntermediateTensors | None,
2202
dict[str, Any],
2203
+ ECConnectorOutput | None,
2204
]:
2205
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2206
is_first_rank = get_pp_group().is_first_rank
2207
2208
# _prepare_inputs may reorder the batch, so we must gather multi
2209
# modal outputs after that to ensure the correct order
2210
+ ec_connector_output = None
2211
+
2212
if (
2213
self.supports_mm_inputs
2214
and is_first_rank
2215
and not self.model_config.is_encoder_decoder
2216
):
2217
# Run the multimodal encoder if any.
2206
- self._execute_mm_encoder(scheduler_output)
2207
- mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
2218
+ with self.maybe_get_ec_connector_output(
2219
+ scheduler_output,
2220
+ encoder_cache=self.encoder_cache,
2221
+ ) as ec_connector_output:
2222
+ self._execute_mm_encoder(scheduler_output)
2223
+ mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
2224
2225
# NOTE(woosuk): To unify token ids and soft tokens (vision
2226
# embeddings), we always use embeddings (rather than token ids)
@@ -2284,6 +2300,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2300
positions,
2301
intermediate_tensors,
2302
model_kwargs,
2303
+ ec_connector_output,
2304
)
2305
2306
def _sample(
@@ -2508,6 +2525,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2525
# Update persistent batch states.
2526
self._update_states(scheduler_output)
2527
2528
+ if has_ec_transfer() and get_ec_transfer().is_producer:
2529
+ with self.maybe_get_ec_connector_output(
2530
+ scheduler_output,
2531
+ encoder_cache=self.encoder_cache,
2532
+ ) as ec_connector_output:
2533
+ self._execute_mm_encoder(scheduler_output)
2534
+ return make_empty_encoder_model_runner_output(scheduler_output)
2535
+
2536
if not num_scheduled_tokens:
2537
if not has_kv_transfer_group():
2538
# Return empty ModelRunnerOutput if no work to do.
@@ -2583,6 +2608,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2608
positions,
2609
intermediate_tensors,
2610
model_kwargs,
2611
+ ec_connector_output,
2612
) = self._preprocess(
2613
scheduler_output, num_input_tokens, intermediate_tensors
2614
)
@@ -2699,6 +2725,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2725
sample_hidden_states,
2726
aux_hidden_states,
2727
kv_connector_output,
2728
+ ec_connector_output,
2729
)
2730
return None
2731
@@ -2720,6 +2747,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2747
sample_hidden_states,
2748
aux_hidden_states,
2749
kv_connector_output,
2750
+ ec_connector_output,
2751
) = self.execute_model_state
2752
# Clear ephemeral state.
2753
self.execute_model_state = None
@@ -2811,6 +2839,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2839
prompt_logprobs_dict=prompt_logprobs_dict,
2840
pooler_output=[],
2841
kv_connector_output=kv_connector_output,
2842
+ ec_connector_output=ec_connector_output
2843
+ if self.supports_mm_inputs
2844
+ else None,
2845
num_nans_in_logits=num_nans_in_logits,
2846
)
2847
@@ -4797,7 +4828,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4828
KVCacheSpec: A dictionary mapping layer names to their KV cache
4829
format. Layers that do not need KV cache are not included.
4830
"""
4800
-
4831
+ if has_ec_transfer() and get_ec_transfer().is_producer:
4832
+ return {}
4833
kv_cache_spec: dict[str, KVCacheSpec] = {}
4834
attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
4835
for layer_name, attn_module in attn_layers.items():
vllm/v1/worker/gpu_worker.py
+5
@@ -20,6 +20,7 @@ from vllm.distributed import (
20
init_distributed_environment,
21
set_custom_all_reduce,
22
)
23
+from vllm.distributed.ec_transfer import ensure_ec_transfer_initialized
24
from vllm.distributed.kv_transfer import (
25
ensure_kv_transfer_initialized,
26
get_kv_transfer_group,
@@ -887,3 +888,7 @@ def init_worker_distributed_environment(
888
parallel_config.pipeline_parallel_size,
889
parallel_config.decode_context_parallel_size,
890
)
891
+
892
+ # Init ec connector here before KV caches caches init
893
+ # NOTE: We do not init KV caches for Encoder-only instance in EPD disagg mode
894
+ ensure_ec_transfer_initialized(vllm_config)