[Model Runner v2] E/P/D disaggregation support (#38390)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Wentao Ye committed
Aug 4, 2026 at 11:52 UTC
4f819f801b7702e39d9588cc9f2ecd28560b31f5
7 files changed
+145
-8
examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh
+1
-1
@@ -29,7 +29,7 @@ if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then
29
fi
30
31
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
32
-TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}" # wait_for_server timeout
32
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-300}" # wait_for_server timeout
33
34
NUM_PROMPTS="${NUM_PROMPTS:-100}" # number of prompts to send in benchmark
35
vllm/config/vllm.py
+18
-4
@@ -536,6 +536,24 @@ class VllmConfig:
536
]
537
return hash_str
538
539
+ @property
540
+ def is_ec_producer_only(self) -> bool:
541
+ ec_config = self.ec_transfer_config
542
+ return (
543
+ ec_config is not None
544
+ and ec_config.is_ec_producer
545
+ and not ec_config.is_ec_consumer
546
+ )
547
+
548
+ @property
549
+ def is_encoder_only(self) -> bool:
550
+ mm_config = (
551
+ self.model_config.multimodal_config
552
+ if self.model_config is not None
553
+ else None
554
+ )
555
+ return self.is_ec_producer_only or bool(mm_config and mm_config.mm_encoder_only)
556
+
557
@property
558
def max_concurrent_batches(self) -> int:
559
# PP requires PP-size concurrent batches to fill the pipeline.
@@ -2246,10 +2264,6 @@ class VllmConfig:
2264
# Will be added by https://github.com/vllm-project/vllm/pull/35045
2265
unsupported.append("KV sharing fast prefill")
2266
2249
- if self.ec_transfer_config is not None:
2250
- # Will be added by https://github.com/vllm-project/vllm/pull/38390
2251
- unsupported.append("EC transfer")
2252
-
2267
return unsupported
2268
2269
def _validate_v2_model_runner(self) -> None:
vllm/v1/worker/gpu/block_table.py
+6
@@ -120,6 +120,8 @@ class BlockTables:
120
self.num_blocks.np[i, req_index] = start + len(block_ids)
121
122
def apply_staged_writes(self) -> None:
123
+ if self.num_kv_cache_groups == 0:
124
+ return
125
if self.num_kv_cache_groups == 1:
126
# Single group: write directly, skipping the per-write group lookup.
127
self.block_tables[0].apply_write()
@@ -138,6 +140,8 @@ class BlockTables:
140
out: tuple[torch.Tensor, ...] | None = None,
141
out_ptrs: torch.Tensor | None = None,
142
) -> tuple[torch.Tensor, ...]:
143
+ if self.num_kv_cache_groups == 0:
144
+ return ()
145
if out is None:
146
out = tuple(self.input_block_tables)
147
out_ptrs = self.input_block_table_ptrs
@@ -179,6 +183,8 @@ class BlockTables:
183
num_tokens_padded: int,
184
out: torch.Tensor | None = None,
185
) -> torch.Tensor:
186
+ if self.num_kv_cache_groups == 0:
187
+ return (self.slot_mappings if out is None else out)[:, :num_tokens_padded]
188
num_reqs = idx_mapping.shape[0]
189
num_groups = self.num_kv_cache_groups
190
slot_mappings = self.slot_mappings if out is None else out
vllm/v1/worker/gpu/ec_connector.py
new
+85
@@ -0,0 +1,85 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Generator
4
+from contextlib import contextmanager
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
11
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase
12
+from vllm.v1.outputs import ECConnectorOutput
13
+
14
+if TYPE_CHECKING:
15
+ from vllm.v1.core.sched.output import SchedulerOutput
16
+ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
17
+
18
+
19
+class ECConnector:
20
+ """EC connector interface used by the V2 GPU model runner."""
21
+
22
+ @contextmanager
23
+ def maybe_get_output(
24
+ self, scheduler_output: "SchedulerOutput"
25
+ ) -> Generator[ECConnectorOutput | None, None, None]:
26
+ yield None
27
+
28
+
29
+class ActiveECConnector(ECConnector):
30
+ def __init__(
31
+ self,
32
+ vllm_config: VllmConfig,
33
+ encoder_cache: dict[str, torch.Tensor],
34
+ ) -> None:
35
+ self.encoder_cache = encoder_cache
36
+ self.save_new_caches = vllm_config.is_ec_producer_only
37
+ self.ec_connector = get_ec_transfer()
38
+ assert isinstance(self.ec_connector, ECConnectorBase)
39
+
40
+ @contextmanager
41
+ def maybe_get_output(
42
+ self, scheduler_output: "SchedulerOutput"
43
+ ) -> Generator[ECConnectorOutput | None, None, None]:
44
+ if scheduler_output.ec_connector_metadata is None:
45
+ yield None
46
+ return
47
+
48
+ output = ECConnectorOutput()
49
+ ec_connector = self.ec_connector
50
+ assert scheduler_output.ec_connector_metadata is not None
51
+ ec_connector.bind_connector_metadata(scheduler_output.ec_connector_metadata)
52
+
53
+ if ec_connector.is_consumer:
54
+ ec_connector.start_load_caches(self.encoder_cache)
55
+
56
+ cached_hashes = set(self.encoder_cache) if self.save_new_caches else None
57
+ try:
58
+ yield output
59
+ if cached_hashes is not None:
60
+ for mm_hash in self.encoder_cache.keys() - cached_hashes:
61
+ ec_connector.save_caches(
62
+ encoder_cache=self.encoder_cache, mm_hash=mm_hash
63
+ )
64
+ finally:
65
+ output.finished_sending, output.finished_recving = (
66
+ ec_connector.get_finished(scheduler_output.finished_req_ids)
67
+ )
68
+ ec_connector.clear_connector_metadata()
69
+
70
+
71
+NO_OP_EC_CONNECTOR = ECConnector()
72
+
73
+
74
+def get_ec_connector(
75
+ vllm_config: VllmConfig,
76
+ encoder_cache: "EncoderCache | None",
77
+) -> ECConnector:
78
+ if (
79
+ not has_ec_transfer()
80
+ or vllm_config.model_config.is_encoder_decoder
81
+ or encoder_cache is None
82
+ ):
83
+ return NO_OP_EC_CONNECTOR
84
+
85
+ return ActiveECConnector(vllm_config, encoder_cache.encoder_outputs)
vllm/v1/worker/gpu/mm/encoder_runner.py
+2
@@ -51,6 +51,8 @@ class EncoderRunner:
51
mm_feature = mm_features[mm_input_id]
52
if mm_feature.data is None:
53
continue
54
+ if mm_feature.identifier in self.encoder_cache.encoder_outputs:
55
+ continue
56
mm_hashes.append(mm_feature.identifier)
57
mm_kwargs.append((mm_feature.modality, mm_feature.data))
58
vllm/v1/worker/gpu/model_runner.py
+30
-3
@@ -63,6 +63,7 @@ from vllm.v1.outputs import (
63
DraftTokenIds,
64
ModelRunnerOutput,
65
RoutedExpertsTensors,
66
+ make_empty_encoder_model_runner_output,
67
)
68
from vllm.v1.worker.block_table import get_block_table_width
69
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
@@ -86,6 +87,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
87
get_uniform_token_count,
88
)
89
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
90
+from vllm.v1.worker.gpu.ec_connector import get_ec_connector
91
from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
92
from vllm.v1.worker.gpu.input_batch import (
93
InputBatch,
@@ -147,6 +149,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
149
150
self.device = device
151
self.dtype = self.model_config.dtype
152
+ self.is_encoder_only = vllm_config.is_encoder_only
153
self.kv_cache_dtype = self.dtype
154
if self.cache_config.cache_dtype != "auto":
155
# Quantized KV cache.
@@ -205,6 +208,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
208
self.encoder_cache = None
209
if self.supports_mm_inputs and self.is_first_pp_rank:
210
self.encoder_cache = EncoderCache()
211
+ self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache)
212
213
# Speculative decoding.
214
self.speculator = None
@@ -448,6 +452,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
452
return torch.cuda.current_stream(self.device)
453
454
def get_kv_cache_spec(self):
455
+ if self.is_encoder_only:
456
+ return {}
457
return get_kv_cache_spec(self.vllm_config)
458
459
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
@@ -580,6 +586,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
586
is_profile: bool = False,
587
**kwargs,
588
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
589
+ if self.is_encoder_only:
590
+ empty = torch.empty(0, device=self.device)
591
+ return empty, empty
592
if skip_attn and not is_profile:
593
raise ValueError(
594
"skip_attn must only be True for initial memory profiling."
@@ -734,6 +743,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
743
dummy_mm_inputs, mm_budget
744
)
745
746
+ if self.is_encoder_only:
747
+ torch.accelerator.synchronize()
748
+ self.reset_encoder_cache()
749
+ gc.collect()
750
+ return
751
+
752
hidden_states, sample_hidden_states = self._dummy_run(
753
self.max_num_tokens, skip_attn=True, is_profile=True
754
)
@@ -770,6 +785,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
785
786
@torch.inference_mode()
787
def capture_model(self) -> int:
788
+ if self.is_encoder_only:
789
+ return 0
790
+
791
assert self.cudagraph_manager is not None
792
if not self.cudagraph_manager.needs_capture():
793
logger.warning(
@@ -1345,6 +1363,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1363
1364
input_ids = input_batch.input_ids
1365
inputs_embeds = None
1366
+ ec_connector_output = None
1367
if self.supports_mm_inputs and self.is_first_pp_rank:
1368
# Run MM encoder (if needed) and get multimodal embeddings.
1369
# Only first PP rank prepares multimodal embeddings.
@@ -1364,12 +1383,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1383
lora_state=self.lora_state,
1384
scheduled_encoder_inputs=scheduled_encoder_inputs,
1385
)
1367
- inputs_embeds = self.model_state.get_mm_embeddings(
1368
- scheduled_encoder_inputs, input_batch, self.req_states
1369
- )
1386
+ with self.ec_connector.maybe_get_output(
1387
+ scheduler_output
1388
+ ) as ec_connector_output:
1389
+ inputs_embeds = self.model_state.get_mm_embeddings(
1390
+ scheduled_encoder_inputs, input_batch, self.req_states
1391
+ )
1392
if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
1393
input_ids = None
1394
1395
+ if self.is_encoder_only:
1396
+ output = make_empty_encoder_model_runner_output(scheduler_output)
1397
+ output.ec_connector_output = ec_connector_output
1398
+ return output
1399
+
1400
model_inputs = {
1401
"input_ids": input_ids,
1402
"positions": input_batch.positions,
vllm/v1/worker/gpu/warmup.py
+3
@@ -162,6 +162,9 @@ def warmup_kernels(
162
We must call the provided worker's execute_model for pipeline parallel
163
coordination.
164
"""
165
+ if model_runner.is_encoder_only:
166
+ return
167
+
168
num_spec_steps = model_runner.num_speculative_steps
169
decode_query_len = model_runner.decode_query_len
170
# Use decode_query_len + 1 tokens so the prefill batch's per-request query