A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
[BugFix] Fix implicit and incorrect assumption on ECConnector is_producer (#34783)
Signed-off-by: Qi Wang <qiwa@nvidia.com>
Qi Wang committed
Mar 4, 2026 at 06:01 UTC
6aa6ad8992a928777f840a843f897ed4cb04c763
8 files changed
+125
-33
examples/online_serving/ec_both_encoder/ec_both_encoder.sh
new
+73
@@ -0,0 +1,73 @@
1
+#!/bin/bash
2
+set -euo pipefail
3
+
4
+MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
5
+PORT="${PORT:-8000}"
6
+GPU="${GPU:-0}"
7
+NUM_PROMPTS="${NUM_PROMPTS:-200}"
8
+EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
9
+TIMEOUT="${TIMEOUT:-600}"
10
+
11
+SERVER_PID=""
12
+
13
+cleanup() {
14
+ echo "Stopping server..."
15
+ if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
16
+ kill "$SERVER_PID" 2>/dev/null || true
17
+ wait "$SERVER_PID" 2>/dev/null || true
18
+ fi
19
+ echo "Done."
20
+}
21
+trap cleanup EXIT INT TERM
22
+
23
+wait_for_server() {
24
+ local deadline=$((SECONDS + TIMEOUT))
25
+ echo "Waiting for server on port $PORT..."
26
+ while (( SECONDS < deadline )); do
27
+ if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then
28
+ echo "Server ready."
29
+ return 0
30
+ fi
31
+ sleep 2
32
+ done
33
+ echo "ERROR: Server did not start within ${TIMEOUT}s"
34
+ return 1
35
+}
36
+
37
+rm -rf "$EC_SHARED_STORAGE_PATH"
38
+mkdir -p "$EC_SHARED_STORAGE_PATH"
39
+
40
+###############################################################################
41
+# Start server with ec_both
42
+###############################################################################
43
+CUDA_VISIBLE_DEVICES="$GPU" \
44
+vllm serve "$MODEL" \
45
+ --port "$PORT" \
46
+ --enforce-eager \
47
+ --ec-transfer-config '{
48
+ "ec_connector": "ECExampleConnector",
49
+ "ec_role": "ec_both",
50
+ "ec_connector_extra_config": {
51
+ "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
52
+ }
53
+ }' \
54
+ "$@" &
55
+
56
+SERVER_PID=$!
57
+wait_for_server
58
+
59
+###############################################################################
60
+# Benchmark -- dataset contains duplicate images, exercises cache hits
61
+###############################################################################
62
+echo "Running benchmark ($NUM_PROMPTS prompts)..."
63
+vllm bench serve \
64
+ --model "$MODEL" \
65
+ --backend openai-chat \
66
+ --endpoint /v1/chat/completions \
67
+ --dataset-name hf \
68
+ --dataset-path lmarena-ai/VisionArena-Chat \
69
+ --seed 0 \
70
+ --num-prompts "$NUM_PROMPTS" \
71
+ --port "$PORT"
72
+
73
+echo "Benchmark complete."
tests/v1/core/test_scheduler.py
+22
-9
@@ -3010,12 +3010,16 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):
3010
# Encoder cache should contain all mm items from request2
3011
_assert_right_encoder_cache_allocated(scheduler, requests=[request2])
3012
3013
- # Should call update_state_after_alloc for hash1_C, ONLY
3013
# hash1_A should not be loaded from connector
3014
# since it's computed in last request & exist in local cache
3015
# Order of getting encoder cache should be: local cache -> connector-> compute
3017
- scheduler.ec_connector.update_state_after_alloc.assert_called_with(request2, 0)
3018
- scheduler.ec_connector.update_state_after_alloc.assert_called_once()
3016
+ # update_state_after_alloc is called for all paths:
3017
+ # index 0 (hash1_C): connector hit → queued for load
3018
+ # index 1 (hash1_D): cache miss → no-op inside connector
3019
+ # index 2 (hash1_E): cache miss → no-op inside connector
3020
+ scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 0)
3021
+ scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 1)
3022
+ scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 2)
3023
3024
scheduler.ec_connector.update_state_after_alloc.reset_mock()
3025
@@ -3087,7 +3091,6 @@ def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):
3091
# mm_hashes of requests exist in cache after scheduling for all scenario
3092
_assert_right_encoder_cache_allocated(scheduler, requests=requests)
3093
3090
- # Should only call update_state_after_alloc when loaded externally
3094
if cache_exist == "connector_only":
3095
scheduler.ec_connector.update_state_after_alloc.assert_called_with(
3096
requests[-1], 0
@@ -3098,9 +3101,15 @@ def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):
3101
3102
# Check metadata should contain mm data for all 10 requests
3103
_assert_right_ec_connector_metadata(output, mm_features_list=mm_features_list)
3101
- else:
3104
+ elif cache_exist == "local":
3105
+ # Local cache hit: items never reach update_state_after_alloc
3106
scheduler.ec_connector.update_state_after_alloc.assert_not_called()
3103
- # ECConnector should carry no metadata
3107
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
3108
+ else:
3109
+ # no_where: called from encoder_inputs_to_schedule but no-op
3110
+ # inside connector (has_cache_item returns False)
3111
+ assert cache_exist == "no_where"
3112
+ scheduler.ec_connector.update_state_after_alloc.assert_called()
3113
_assert_right_ec_connector_metadata(output, mm_features_list=[])
3114
3115
scheduler.ec_connector.update_state_after_alloc.reset_mock()
@@ -3419,7 +3428,6 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
3428
# mm_hash of request_low exists in cache after scheduling for all scenario
3429
_assert_right_encoder_cache_allocated(scheduler, requests=[request_low])
3430
3422
- # Should only call update_state_after_alloc when loaded externally
3431
if cache_exist == "connector_only":
3432
scheduler.ec_connector.update_state_after_alloc.assert_called_with(
3433
request_low, 0
@@ -3427,9 +3435,14 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
3435
_assert_right_ec_connector_metadata(
3436
output, mm_features_list=request_low.mm_features
3437
)
3430
- else:
3438
+ elif cache_exist == "local":
3439
scheduler.ec_connector.update_state_after_alloc.assert_not_called()
3432
- # ECConnector should carry no metadata
3440
+ _assert_right_ec_connector_metadata(output, mm_features_list=[])
3441
+ else:
3442
+ assert cache_exist == "no_where"
3443
+ scheduler.ec_connector.update_state_after_alloc.assert_called_with(
3444
+ request_low, 0
3445
+ )
3446
_assert_right_ec_connector_metadata(output, mm_features_list=[])
3447
3448
scheduler.ec_connector.update_state_after_alloc.reset_mock()
tests/v1/ec_connector/unit/test_ec_example_connector.py
+16
-16
@@ -233,9 +233,10 @@ class TestStateManagement:
233
# Initial state should be empty
234
assert len(connector._mm_datas_need_loads) == 0
235
236
- # Update state for all 3 items
237
- for i in range(3):
238
- connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
236
+ # Update state for all 3 items (mock cache existence)
237
+ with patch.object(connector, "has_cache_item", return_value=True):
238
+ for i in range(3):
239
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
240
241
# Check state updated for all 3
242
assert len(connector._mm_datas_need_loads) == 3
@@ -255,9 +256,10 @@ class TestStateManagement:
256
role=ECConnectorRole.SCHEDULER,
257
)
258
258
- # Setup state for all 3 items
259
- for i in range(3):
260
- connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
259
+ # Setup state for all 3 items (mock cache existence)
260
+ with patch.object(connector, "has_cache_item", return_value=True):
261
+ for i in range(3):
262
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
263
264
# Build metadata
265
scheduler_output = Mock(spec=SchedulerOutput)
@@ -298,9 +300,10 @@ class TestStateManagement:
300
role=ECConnectorRole.SCHEDULER,
301
)
302
301
- # Add state
302
- for i in range(3):
303
- connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
303
+ # Add state (mock cache existence)
304
+ with patch.object(connector, "has_cache_item", return_value=True):
305
+ for i in range(3):
306
+ connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
307
assert len(connector._mm_datas_need_loads) == 3
308
309
# Build metadata (should clear state)
@@ -608,16 +611,13 @@ class TestEdgeCases:
611
with pytest.raises(FileNotFoundError):
612
connector.start_load_caches(encoder_cache=encoder_cache)
613
611
- def test_has_caches_empty_request(self, mock_vllm_config_producer):
612
- """Test has_caches with request that has no MM data."""
614
+ def test_has_cache_item_empty_request(self, mock_vllm_config_producer):
615
+ """Test has_cache_item with a nonexistent identifier."""
616
connector = ECExampleConnector(
617
vllm_config=mock_vllm_config_producer,
618
role=ECConnectorRole.SCHEDULER,
619
)
620
618
- mock_request = MockRequest("req_empty", [], [])
621
+ result = connector.has_cache_item("nonexistent_hash")
622
620
- result = connector.has_caches(mock_request)
621
-
622
- assert len(result) == 0
623
- assert result == []
623
+ assert result is False
vllm/distributed/ec_transfer/ec_connector/example_connector.py
+3
-1
@@ -141,8 +141,10 @@ class ECExampleConnector(ECConnectorBase):
141
Update ECConnector state after encoder cache allocation.
142
"""
143
mm_hash = request.mm_features[index].identifier
144
+ # Only load cache if it is consumer and cache exists
145
+ if not self.is_consumer or not self.has_cache_item(mm_hash):
146
+ return
147
num_encoder_token = request.get_num_encoder_embeds(index)
145
- # Insert mm_hash only if this block has not been recorded yet.
148
self._mm_datas_need_loads[mm_hash] = num_encoder_token
149
150
def build_connector_meta(
vllm/v1/core/sched/scheduler.py
+4
@@ -515,6 +515,8 @@ class Scheduler(SchedulerInterface):
515
# Allocate the encoder cache.
516
for i in encoder_inputs_to_schedule:
517
self.encoder_cache_manager.allocate(request, i)
518
+ if self.ec_connector is not None:
519
+ self.ec_connector.update_state_after_alloc(request, i)
520
encoder_compute_budget = new_encoder_compute_budget
521
if external_load_encoder_input:
522
for i in external_load_encoder_input:
@@ -803,6 +805,8 @@ class Scheduler(SchedulerInterface):
805
# Allocate the encoder cache.
806
for i in encoder_inputs_to_schedule:
807
self.encoder_cache_manager.allocate(request, i)
808
+ if self.ec_connector is not None:
809
+ self.ec_connector.update_state_after_alloc(request, i)
810
encoder_compute_budget = new_encoder_compute_budget
811
# Allocate for external load encoder cache
812
if external_load_encoder_input:
vllm/v1/engine/core.py
+4
-4
@@ -193,9 +193,9 @@ class EngineCore:
193
logger.debug("Batch queue is enabled with size %d", self.batch_queue_size)
194
self.batch_queue = deque(maxlen=self.batch_queue_size)
195
196
- self.is_ec_producer = (
197
- vllm_config.ec_transfer_config is not None
198
- and vllm_config.ec_transfer_config.is_ec_producer
196
+ self.is_ec_consumer = (
197
+ vllm_config.ec_transfer_config is None
198
+ or vllm_config.ec_transfer_config.is_ec_consumer
199
)
200
self.is_pooling_model = vllm_config.model_config.runner_type == "pooling"
201
@@ -449,7 +449,7 @@ class EngineCore:
449
exec_future = self.model_executor.execute_model(
450
scheduler_output, non_block=True
451
)
452
- if not self.is_ec_producer:
452
+ if self.is_ec_consumer:
453
model_executed = scheduler_output.total_num_scheduled_tokens > 0
454
455
if self.is_pooling_model or not model_executed:
vllm/v1/executor/ray_executor.py
+1
-1
@@ -100,7 +100,7 @@ class RayDistributedExecutor(Executor):
100
101
self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
102
self.vllm_config.ec_transfer_config is None
103
- or not self.vllm_config.ec_transfer_config.is_ec_producer
103
+ or self.vllm_config.ec_transfer_config.is_ec_consumer
104
)
105
106
self.scheduler_output: SchedulerOutput | None = None
vllm/v1/worker/gpu_model_runner.py
+2
-2
@@ -3409,7 +3409,7 @@ class GPUModelRunner(
3409
# Update persistent batch states.
3410
self._update_states(scheduler_output)
3411
3412
- if has_ec_transfer() and get_ec_transfer().is_producer:
3412
+ if has_ec_transfer() and not get_ec_transfer().is_consumer:
3413
with self.maybe_get_ec_connector_output(
3414
scheduler_output,
3415
encoder_cache=self.encoder_cache,
@@ -6182,7 +6182,7 @@ class GPUModelRunner(
6182
KVCacheSpec: A dictionary mapping layer names to their KV cache
6183
format. Layers that do not need KV cache are not included.
6184
"""
6185
- if has_ec_transfer() and get_ec_transfer().is_producer:
6185
+ if has_ec_transfer() and not get_ec_transfer().is_consumer:
6186
return {}
6187
kv_cache_spec: dict[str, KVCacheSpec] = {}
6188
layer_type = cast(type[Any], AttentionLayerBase)