@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fa9d67f782

[EC Connector] Added Build Connector Worker Meta for EC Connector (#49585)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

omerpaz95 committed Aug 16, 2026 at 01:40 UTC fa9d67f7828e9bc105912ddf41dc384105732b1e
16 files changed +430 -16
.buildkite/test_areas/misc.yaml
+1
@@ -127,6 +127,7 @@ steps:
127 - pytest -v -s v1/test_kv_cache_spec_registry.py
128 - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
129 - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
130 + - pytest -v -s -m 'cpu_test' v1/ec_connector/unit
131 - pytest -v -s -m 'cpu_test' v1/metrics
132
133 - label: Extract Hidden States Integration
tests/v1/ec_connector/unit/test_ec_output_aggregator.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ECOutputAggregator."""
4 +
5 +import pytest
6 +
7 +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
8 +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
9 +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
10 +from vllm.v1.outputs import (
11 + EMPTY_MODEL_RUNNER_OUTPUT,
12 + ECConnectorOutput,
13 + KVConnectorOutput,
14 + ModelRunnerOutput,
15 +)
16 +
17 +pytestmark = pytest.mark.cpu_test
18 +
19 +
20 +class FakeWorkerMeta(ECConnectorWorkerMetadata):
21 + """Records merge order. `aggregate` returns a new object, as the base class
22 + declares: an aggregator discarding the return value would lose the merge.
23 + """
24 +
25 + def __init__(self, saves: list[str]):
26 + self.saves = saves
27 +
28 + def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta":
29 + return FakeWorkerMeta(self.saves + other.saves)
30 +
31 +
32 +def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput:
33 + return ModelRunnerOutput(
34 + req_ids=[], req_id_to_index={}, ec_connector_output=ec_output
35 + )
36 +
37 +
38 +def test_aggregate_folds_every_rank_onto_output_rank():
39 + """EC work done on any rank reaches the scheduler via output_rank's output.
40 +
41 + The middle rank reports no worker metadata: it must neither seed nor clobber
42 + the accumulator.
43 + """
44 + outputs = [
45 + _worker_output(
46 + ECConnectorOutput(
47 + finished_sending={"mm0"},
48 + ec_connector_worker_meta=FakeWorkerMeta(["mm0"]),
49 + )
50 + ),
51 + _worker_output(ECConnectorOutput(finished_recving={"mm1"})),
52 + _worker_output(
53 + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"]))
54 + ),
55 + ]
56 +
57 + result = ECOutputAggregator().aggregate(outputs, output_rank=2)
58 +
59 + assert result is outputs[2]
60 + assert result.ec_connector_output.finished_sending == {"mm0"}
61 + assert result.ec_connector_output.finished_recving == {"mm1"}
62 + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"]
63 +
64 +
65 +def test_aggregate_leaves_no_ec_output_when_no_worker_reported():
66 + """Empty per-worker reports must not reach the scheduler as an empty object."""
67 + outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())]
68 +
69 + result = ECOutputAggregator().aggregate(outputs, output_rank=0)
70 +
71 + assert result is outputs[0]
72 + assert result.ec_connector_output is None
73 + assert ECOutputAggregator().aggregate([None], output_rank=0) is None
74 +
75 +
76 +def test_aggregate_does_not_write_through_the_shared_empty_output():
77 + """A rank with nothing to report yields the shared empty output singleton.
78 +
79 + Folding another rank's metadata onto it must not write through to the
80 + module-level object, which every later step would then carry.
81 + """
82 + outputs = [
83 + _worker_output(
84 + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"]))
85 + ),
86 + EMPTY_MODEL_RUNNER_OUTPUT,
87 + ]
88 +
89 + result = ECOutputAggregator().aggregate(outputs, output_rank=1)
90 +
91 + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None
92 + assert result is not EMPTY_MODEL_RUNNER_OUTPUT
93 + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"]
94 +
95 +
96 +def test_chaining_with_kv_aggregator_preserves_both_outputs():
97 + """MultiprocExecutor chains both aggregators and keeps only the last result,
98 + so each must merge onto the same output_rank output rather than replace it.
99 + """
100 + outputs = [
101 + _worker_output(ECConnectorOutput(finished_sending={"mm0"})),
102 + _worker_output(None),
103 + ]
104 + outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7})
105 +
106 + result = None
107 + for aggregator in (
108 + KVOutputAggregator(expected_finished_count=1),
109 + ECOutputAggregator(),
110 + ):
111 + result = aggregator.aggregate(outputs, output_rank=1)
112 +
113 + assert result is outputs[1]
114 + assert result.kv_connector_output.invalid_block_ids == {7}
115 + assert result.ec_connector_output.finished_sending == {"mm0"}
tests/v1/ec_connector/unit/test_worker_ec_connector.py new
+77
@@ -0,0 +1,77 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the V2 GPU model runner's EC connector wrapper."""
4 +
5 +from types import SimpleNamespace
6 +from unittest.mock import MagicMock, patch
7 +
8 +import pytest
9 +
10 +from vllm.distributed.ec_transfer.ec_connector.base import (
11 + ECConnectorBase,
12 + ECConnectorMetadata,
13 +)
14 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
15 +from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector
16 +
17 +pytestmark = pytest.mark.cpu_test
18 +
19 +WORKER_META = object()
20 +
21 +
22 +def _scheduler_output() -> SimpleNamespace:
23 + return SimpleNamespace(
24 + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()
25 + )
26 +
27 +
28 +def _connector(
29 + encoder_cache: dict | None = None,
30 + is_producer: bool = True,
31 + is_consumer: bool = False,
32 +) -> tuple[ActiveECConnector, MagicMock]:
33 + fake = MagicMock(spec=ECConnectorBase)
34 + fake.is_producer = is_producer
35 + fake.is_consumer = is_consumer
36 + fake.get_finished.return_value = (None, None)
37 + fake.build_connector_worker_meta.return_value = WORKER_META
38 + with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake):
39 + return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake
40 +
41 +
42 +@pytest.mark.parametrize(
43 + ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)]
44 +)
45 +def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer):
46 + """An ec_both node is also a producer: it must offload what it just computed."""
47 + encoder_cache = {"mm_old": None}
48 + connector, fake = _connector(encoder_cache, is_producer, is_consumer)
49 +
50 + with connector.maybe_get_output(_scheduler_output()):
51 + encoder_cache["mm_new"] = None
52 +
53 + saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list]
54 + assert saved == (["mm_new"] if is_producer else [])
55 + assert fake.start_load_caches.called == is_consumer
56 +
57 +
58 +def test_worker_meta_is_reported_on_context_exit():
59 + """Reported in the finally block, so is_empty() sees it only after the exit."""
60 + connector, fake = _connector()
61 +
62 + with connector.maybe_get_output(_scheduler_output()) as output:
63 + assert output.ec_connector_worker_meta is None
64 +
65 + assert output.ec_connector_worker_meta is WORKER_META
66 + assert fake.clear_connector_metadata.called
67 +
68 +
69 +def test_no_forward_reports_without_running_the_model():
70 + connector, _ = _connector()
71 +
72 + output = connector.no_forward(_scheduler_output())
73 +
74 + assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META
75 +
76 + empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output())
77 + assert empty is EMPTY_MODEL_RUNNER_OUTPUT
tests/v1/executor/test_executor.py
+3
@@ -9,6 +9,7 @@ from typing import Any
9
10 import pytest
11
12 +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
13 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
14 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
15 from vllm.sampling_params import SamplingParams
@@ -98,6 +99,7 @@ class CustomMultiprocExecutor(MultiprocExecutor):
99 non_block: bool = False,
100 unique_reply_rank: int | None = None,
101 kv_output_aggregator: KVOutputAggregator = None,
102 + ec_output_aggregator: ECOutputAggregator | None = None,
103 ) -> Any | list[Any] | Future[Any | list[Any]]:
104 # Drop marker to show that this was run
105 with open(".marker", "w"):
@@ -110,6 +112,7 @@ class CustomMultiprocExecutor(MultiprocExecutor):
112 non_block,
113 unique_reply_rank,
114 kv_output_aggregator,
115 + ec_output_aggregator,
116 )
117
118
tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
+1
@@ -55,6 +55,7 @@ def _run_engine_core_handshake(
55 class _FakeScheduler:
56 def __init__(self, **kwargs: Any) -> None:
57 self.connector = connector
58 + self.ec_connector = None
59
60 def get_kv_connector(self) -> KVConnectorBase_V1:
61 return connector
tests/v1/test_outputs.py
+18 -1
@@ -6,7 +6,13 @@ import numpy as np
6 import torch
7
8 from vllm.platforms import current_platform
9 -from vllm.v1.outputs import LogprobsLists, LogprobsTensors
9 +from vllm.v1.outputs import (
10 + EMPTY_MODEL_RUNNER_OUTPUT,
11 + ECConnectorOutput,
12 + LogprobsLists,
13 + LogprobsTensors,
14 + ModelRunnerOutput,
15 +)
16 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
17 from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
18
@@ -200,3 +206,14 @@ class TestLogprobsLists(TestCase):
206 assert len(sliced.logprob_token_ids) == 9 # All tokens
207 assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids
208 assert sliced.cu_num_generated_tokens is None
209 +
210 +
211 +def test_with_ec_conn_output_copies_shared_empty_output():
212 + """The shared empty output is copied, never written to."""
213 + ec_output = ECConnectorOutput(finished_sending={"mm_hash"})
214 +
215 + result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output)
216 +
217 + assert result is not EMPTY_MODEL_RUNNER_OUTPUT
218 + assert result.ec_connector_output is ec_output
219 + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
@@ -180,6 +180,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
180 hidden_states=None,
181 aux_hidden_states=None,
182 finished_req_ids=set(),
183 + ec_connector_output=None,
184 routed_experts=None,
185 num_tokens_across_dp=None,
186 )
vllm/distributed/ec_transfer/ec_connector/base.py
+33
@@ -20,6 +20,8 @@ The class provides the following primitives:
20
21 get_finished() - called with ids of finished requests, returns
22 ids of requests that have completed async sending/recving.
23 + build_connector_worker_meta() - builds metadata to be sent
24 + back to the scheduler-side connector
25 """
26
27 import enum
@@ -56,6 +58,27 @@ class ECConnectorMetadata(ABC): # noqa: B024
58 pass
59
60
61 +class ECConnectorWorkerMetadata(ABC):
62 + """
63 + Abstract Metadata used to communicate back
64 + Worker ECConnector -> Scheduler ECConnector.
65 +
66 + Each worker can output its own metadata.
67 + For a single engine step, all metadata objects returned by workers
68 + will be aggregated using the `aggregate` method below, before
69 + being passed to the Scheduler ECConnector.
70 + """
71 +
72 + @abstractmethod
73 + def aggregate(
74 + self, other: "ECConnectorWorkerMetadata"
75 + ) -> "ECConnectorWorkerMetadata":
76 + """
77 + Aggregate metadata with another `ECConnectorWorkerMetadata` object.
78 + """
79 + pass
80 +
81 +
82 class ECConnectorBase(ABC):
83 def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
84 self._connector_metadata: ECConnectorMetadata | None = None
@@ -190,6 +213,16 @@ class ECConnectorBase(ABC):
213 """
214 return None, None
215
216 + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
217 + """
218 + Build the ECConnector worker metadata for this engine step.
219 +
220 + Returns:
221 + ECConnectorWorkerMetadata: the worker metadata.
222 + None if no worker metadata is available.
223 + """
224 + return None
225 +
226 # ==============================
227 # Scheduler-side methods
228 # ==============================
vllm/distributed/ec_transfer/ec_connector/utils.py new
+51
@@ -0,0 +1,51 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""EC connector helper utilities."""
4 +
5 +from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
6 +
7 +
8 +class ECOutputAggregator:
9 + """Merge every worker's EC connector output onto the single
10 + ModelRunnerOutput that reaches the scheduler.
11 +
12 + Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the
13 + scheduler, but the EC connector may have run on any rank.
14 + """
15 +
16 + def aggregate(
17 + self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
18 + ) -> ModelRunnerOutput | None:
19 + output = outputs[output_rank]
20 + if not output:
21 + return None
22 +
23 + finished_sending = set[str]()
24 + finished_recving = set[str]()
25 + worker_meta = None
26 + for model_runner_output in outputs:
27 + assert model_runner_output is not None
28 + ec_output = model_runner_output.ec_connector_output
29 + if not ec_output:
30 + continue
31 +
32 + finished_sending |= ec_output.finished_sending or set()
33 + finished_recving |= ec_output.finished_recving or set()
34 +
35 + if meta := ec_output.ec_connector_worker_meta:
36 + worker_meta = (
37 + meta if worker_meta is None else worker_meta.aggregate(meta)
38 + )
39 +
40 + aggregated = ECConnectorOutput(
41 + finished_sending=finished_sending or None,
42 + finished_recving=finished_recving or None,
43 + ec_connector_worker_meta=worker_meta,
44 + )
45 + if aggregated.is_empty():
46 + output.ec_connector_output = None
47 + return output
48 +
49 + # `output` is the shared empty output whenever `output_rank` had no work,
50 + # so attach through the copy-on-write helper.
51 + return ModelRunnerOutput.with_ec_conn_output(output, aggregated)
vllm/v1/engine/core.py
+2
@@ -173,6 +173,8 @@ class EngineCore:
173 )
174 if self.scheduler.connector is not None: # type: ignore
175 self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore
176 + if self.scheduler.ec_connector is not None: # type: ignore
177 + self.model_executor.init_ec_output_aggregator()
178
179 mm_registry = MULTIMODAL_REGISTRY
180 self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config(
vllm/v1/executor/abstract.py
+5
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Literal, TypeVar, overload
9
10 import vllm.envs as envs
11 from vllm.config import VllmConfig
12 +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
13 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
14 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15 KVConnectorHandshakeMetadata,
@@ -110,6 +111,7 @@ class Executor(ABC):
111 self.is_sleeping = False
112 self.sleeping_tags: set[str] = set()
113 self.kv_output_aggregator: KVOutputAggregator | None = None
114 + self.ec_output_aggregator: ECOutputAggregator | None = None
115
116 @abstractmethod
117 def _init_executor(self) -> None:
@@ -283,6 +285,9 @@ class Executor(ABC):
285 connector, self.parallel_config.world_size
286 )
287
288 + def init_ec_output_aggregator(self) -> None:
289 + self.ec_output_aggregator = ECOutputAggregator()
290 +
291 @cached_property # Avoid unnecessary RPC calls
292 def supported_tasks(self) -> tuple[SupportedTask, ...]:
293 output: list[tuple[SupportedTask, ...]]
vllm/v1/executor/multiproc_executor.py
+20 -6
@@ -29,6 +29,7 @@ import vllm.envs as envs
29 from vllm.config import VllmConfig
30 from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
31 from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue
32 +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
33 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
34 from vllm.distributed.parallel_state import (
35 get_dcp_group,
@@ -343,6 +344,7 @@ class MultiprocExecutor(Executor):
344 non_block=non_block,
345 timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
346 kv_output_aggregator=self.kv_output_aggregator,
347 + ec_output_aggregator=self.ec_output_aggregator,
348 )
349
350 def sample_tokens( # type: ignore[override]
@@ -355,6 +357,7 @@ class MultiprocExecutor(Executor):
357 non_block=non_block,
358 timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
359 kv_output_aggregator=self.kv_output_aggregator,
360 + ec_output_aggregator=self.ec_output_aggregator,
361 )
362
363 def execute_dummy_batch(self) -> None:
@@ -375,9 +378,10 @@ class MultiprocExecutor(Executor):
378 non_block: bool = False,
379 unique_reply_rank: int | None = None,
380 kv_output_aggregator: KVOutputAggregator | None = None,
381 + ec_output_aggregator: ECOutputAggregator | None = None,
382 ) -> Any:
379 - """Returns single result if unique_reply_rank and/or kv_output_aggregator
380 - is provided, otherwise list."""
383 + """Returns single result if unique_reply_rank and/or an output
384 + aggregator is provided, otherwise list."""
385 assert self.rpc_broadcast_mq is not None, (
386 "collective_rpc should not be called on follower node"
387 )
@@ -387,11 +391,21 @@ class MultiprocExecutor(Executor):
391 deadline = None if timeout is None else time.monotonic() + timeout
392 kwargs = kwargs or {}
393
390 - if kv_output_aggregator is not None:
394 + aggregators = [a for a in (kv_output_aggregator, ec_output_aggregator) if a]
395 + aggregate: Callable[[Any], Any]
396 + if aggregators:
397 output_rank = None
392 - aggregate: Callable[[Any], Any] = partial(
393 - kv_output_aggregator.aggregate, output_rank=unique_reply_rank or 0
394 - )
398 +
399 + def _aggregate(outputs: Any) -> Any:
400 + # Each aggregator merges its own connector's output onto
401 + # outputs[output_rank] in place and returns it, so chaining is safe.
402 + rank = unique_reply_rank or 0
403 + result = outputs[rank]
404 + for a in aggregators:
405 + result = a.aggregate(outputs, output_rank=rank)
406 + return result
407 +
408 + aggregate = _aggregate
409 else:
410 output_rank = unique_reply_rank
411 aggregate = lambda x: x
vllm/v1/executor/ray_executor.py
+13
@@ -89,6 +89,19 @@ class RayDistributedExecutor(Executor):
89 # KV connector setup
90 self.has_connector = self.vllm_config.kv_transfer_config is not None
91
92 + if (
93 + self.vllm_config.ec_transfer_config is not None
94 + and self.parallel_config.world_size > 1
95 + ):
96 + raise NotImplementedError(
97 + "EC connector worker metadata is not supported with the "
98 + "legacy Ray executor when world_size > 1: only the output "
99 + "of a single worker is fetched, silently dropping the "
100 + "other workers' EC connector state. Set "
101 + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, "
102 + "or use the multiprocessing executor instead."
103 + )
104 +
105 self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
106 self.vllm_config.ec_transfer_config is None
107 or self.vllm_config.ec_transfer_config.is_ec_consumer
vllm/v1/outputs.py
+36
@@ -14,6 +14,7 @@ from vllm.compilation.cuda_graph import CUDAGraphStat
14 from vllm.v1.core.sched.output import SchedulerOutput
15
16 if TYPE_CHECKING:
17 + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
18 from vllm.distributed.kv_events import KVConnectorKVEvents
19 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
20 KVConnectorWorkerMetadata,
@@ -23,6 +24,7 @@ else:
24 KVConnectorStats = object
25 KVConnectorWorkerMetadata = object
26 KVConnectorKVEvents = object
27 + ECConnectorWorkerMetadata = object
28
29
30 class LogprobsLists(NamedTuple):
@@ -292,6 +294,14 @@ class ECConnectorOutput:
294 # [mm_hash]
295 finished_sending: set[str] | None = None
296 finished_recving: set[str] | None = None
297 + ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None
298 +
299 + def is_empty(self):
300 + return (
301 + not self.finished_sending
302 + and not self.finished_recving
303 + and not self.ec_connector_worker_meta
304 + )
305
306
307 # ModelRunnerOutput is serialized and sent to the scheduler process.
@@ -362,6 +372,32 @@ class ModelRunnerOutput:
372 output.kv_connector_output = kv_connector_output
373 return output
374
375 + @staticmethod
376 + def with_ec_conn_output_only(
377 + ec_connector_output: ECConnectorOutput | None,
378 + ) -> "ModelRunnerOutput":
379 + """Return an otherwise-empty output carrying `ec_connector_output`."""
380 + return ModelRunnerOutput.with_ec_conn_output(
381 + EMPTY_MODEL_RUNNER_OUTPUT, ec_connector_output
382 + )
383 +
384 + @staticmethod
385 + def with_ec_conn_output(
386 + output: "ModelRunnerOutput",
387 + ec_connector_output: ECConnectorOutput | None,
388 + ) -> "ModelRunnerOutput":
389 + """Return `output` carrying `ec_connector_output`.
390 +
391 + The shared empty output is copied rather than written to, so callers
392 + must use the return value.
393 + """
394 + if ec_connector_output is None or ec_connector_output.is_empty():
395 + return output
396 + if output is EMPTY_MODEL_RUNNER_OUTPUT:
397 + output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
398 + output.ec_connector_output = ec_connector_output
399 + return output
400 +
401
402 # ModelRunnerOutput wrapper for async scheduling.
403 class AsyncModelRunnerOutput(ABC):
vllm/v1/worker/gpu/ec_connector.py
+25 -2
@@ -9,7 +9,11 @@ import torch
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
12 +from vllm.v1.outputs import (
13 + EMPTY_MODEL_RUNNER_OUTPUT,
14 + ECConnectorOutput,
15 + ModelRunnerOutput,
16 +)
17
18 if TYPE_CHECKING:
19 from vllm.v1.core.sched.output import SchedulerOutput
@@ -25,6 +29,12 @@ class ECConnector:
29 ) -> Generator[ECConnectorOutput | None, None, None]:
30 yield None
31
32 + def no_forward(
33 + self,
34 + scheduler_output: "SchedulerOutput",
35 + ) -> ModelRunnerOutput:
36 + return EMPTY_MODEL_RUNNER_OUTPUT
37 +
38
39 class ActiveECConnector(ECConnector):
40 def __init__(
@@ -33,9 +43,11 @@ class ActiveECConnector(ECConnector):
43 encoder_cache: dict[str, torch.Tensor],
44 ) -> None:
45 self.encoder_cache = encoder_cache
36 - self.save_new_caches = vllm_config.is_ec_producer_only
46 self.ec_connector = get_ec_transfer()
47 assert isinstance(self.ec_connector, ECConnectorBase)
48 + # Every producer offloads freshly computed encoder outputs, including
49 + # an ec_both node that also reloads them.
50 + self.save_new_caches = self.ec_connector.is_producer
51
52 @contextmanager
53 def maybe_get_output(
@@ -65,8 +77,19 @@ class ActiveECConnector(ECConnector):
77 output.finished_sending, output.finished_recving = (
78 ec_connector.get_finished(scheduler_output.finished_req_ids)
79 )
80 + output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
81 ec_connector.clear_connector_metadata()
82
83 + def no_forward(
84 + self,
85 + scheduler_output: "SchedulerOutput",
86 + ) -> ModelRunnerOutput:
87 + # EC send/recv even if no work to do.
88 + with self.maybe_get_output(scheduler_output) as ec_connector_output:
89 + pass
90 +
91 + return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output)
92 +
93
94 NO_OP_EC_CONNECTOR = ECConnector()
95
vllm/v1/worker/gpu/model_runner.py
+29 -7
@@ -65,6 +65,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
65 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
66 from vllm.v1.outputs import (
67 DraftTokenIds,
68 + ECConnectorOutput,
69 ModelRunnerOutput,
70 RoutedExpertsTensors,
71 make_empty_encoder_model_runner_output,
@@ -1393,6 +1394,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1394 idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu
1395 )
1396
1397 + def _merge_ec_connector_no_forward(
1398 + self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput
1399 + ) -> ModelRunnerOutput:
1400 + """Let the EC connector send/recv on a step with no work to run."""
1401 + return ModelRunnerOutput.with_ec_conn_output(
1402 + output,
1403 + self.ec_connector.no_forward(scheduler_output).ec_connector_output,
1404 + )
1405 +
1406 @torch.inference_mode()
1407 def execute_model(
1408 self,
@@ -1414,7 +1424,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1424 if scheduler_output.total_num_scheduled_tokens == 0:
1425 # No need to run the model.
1426 empty_output = self.kv_connector.no_forward(scheduler_output)
1417 - return empty_output
1427 + return self._merge_ec_connector_no_forward(
1428 + scheduler_output, empty_output
1429 + )
1430
1431 # Get batch descriptor and sync across DP ranks.
1432 num_reqs = len(scheduler_output.num_scheduled_tokens)
@@ -1455,7 +1467,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1467 if batch_desc.num_tokens == 0:
1468 # All DP ranks have zero tokens to run.
1469 empty_output = self.kv_connector.no_forward(scheduler_output)
1458 - return empty_output
1470 + return self._merge_ec_connector_no_forward(scheduler_output, empty_output)
1471
1472 if not dummy_run:
1473 # Common case.
@@ -1572,9 +1584,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1584 input_ids = None
1585
1586 if self.is_encoder_only:
1575 - output = make_empty_encoder_model_runner_output(scheduler_output)
1576 - output.ec_connector_output = ec_connector_output
1577 - return output
1587 + return ModelRunnerOutput.with_ec_conn_output(
1588 + make_empty_encoder_model_runner_output(scheduler_output),
1589 + ec_connector_output,
1590 + )
1591
1592 model_inputs = {
1593 "input_ids": input_ids,
@@ -1679,6 +1692,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1692 hidden_states=hidden_states,
1693 aux_hidden_states=aux_hidden_states,
1694 finished_req_ids=finished_req_ids,
1695 + ec_connector_output=ec_connector_output,
1696 routed_experts=routed_experts,
1697 )
1698
@@ -1702,6 +1716,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1716 hidden_states = self.execute_model_state.hidden_states
1717 aux_hidden_states = self.execute_model_state.aux_hidden_states
1718 finished_req_ids = self.execute_model_state.finished_req_ids
1719 + ec_connector_output = self.execute_model_state.ec_connector_output
1720 routed_experts = self.execute_model_state.routed_experts
1721 self.execute_model_state = None
1722
@@ -1721,7 +1736,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1736
1737 # Post-step KV connector related operations.
1738 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1724 - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1739 + # The first PP rank holds the encoder cache, so pass its EC output on.
1740 + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1741 + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
1742
1743 # Last rank: sample tokens
1744 hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
@@ -1838,6 +1855,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1855 # Post-step KV connector related operations.
1856 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1857 model_runner_output.kv_connector_output = kv_connector_output
1858 + model_runner_output.ec_connector_output = ec_connector_output
1859
1860 return async_output
1861
@@ -1854,6 +1872,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1872 input_batch = self.execute_model_state.input_batch
1873 hidden_states = self.execute_model_state.hidden_states
1874 finished_req_ids = self.execute_model_state.finished_req_ids
1875 + ec_connector_output = self.execute_model_state.ec_connector_output
1876 self.execute_model_state = None
1877
1878 # Post-step KV connector related operations.
@@ -1861,7 +1880,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1880
1881 if not self.is_last_pp_rank:
1882 self.postprocess_num_computed_tokens(input_batch)
1864 - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1883 + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
1884 + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
1885
1886 assert self.pooling_runner is not None
1887 pooler_output, finished_mask = self.pooling_runner.pool(
@@ -1873,6 +1893,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1893 req_ids=input_batch.req_ids,
1894 req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
1895 kv_connector_output=kv_connector_output,
1896 + ec_connector_output=ec_connector_output,
1897 )
1898 async_output = AsyncPoolingOutput(
1899 model_runner_output=model_runner_output,
@@ -1962,6 +1983,7 @@ class ExecuteModelState(NamedTuple):
1983 hidden_states: torch.Tensor | None
1984 aux_hidden_states: list[torch.Tensor] | None
1985 finished_req_ids: set[str]
1986 + ec_connector_output: ECConnectorOutput | None
1987 routed_experts: RoutedExpertsTensors | None
1988
1989