From 2c2c65e974b00032d7e3152f665d0f6f769a4109 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Fri, 27 Mar 2026 19:19:00 +0000 Subject: [PATCH 01/31] e/p/d disaggregation support for MRv2 Signed-off-by: yewentao256 --- vllm/v1/worker/gpu/mm/encoder_runner.py | 2 + vllm/v1/worker/gpu/model_runner.py | 103 +++++++++++++++++++----- vllm/v1/worker/gpu/warmup.py | 3 + 3 files changed, 90 insertions(+), 18 deletions(-) diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index fb2a21ce43e6..ffb0392ed0c0 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -42,6 +42,8 @@ def prepare_mm_inputs( mm_feature = mm_features[mm_input_id] if mm_feature.data is None: continue + if mm_feature.identifier in self.encoder_cache.encoder_outputs: + continue mm_hashes.append(mm_feature.identifier) mm_kwargs.append((mm_feature.modality, mm_feature.data)) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a2f83c52e951..7de08c075b7a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -20,6 +20,7 @@ import functools import gc import time +from contextlib import nullcontext from copy import deepcopy from typing import Any, NamedTuple @@ -44,8 +45,16 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.outputs import ( + DraftTokenIds, + KVConnectorOutput, + ModelRunnerOutput, + make_empty_encoder_model_runner_output, +) from vllm.v1.worker.cp_utils import check_attention_cp_compatibility +from vllm.v1.worker.ec_connector_model_runner_mixin import ( + ECConnectorModelRunnerMixin, +) from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput from vllm.v1.worker.gpu.attn_utils import ( build_slot_mappings_by_layer, @@ -115,6 +124,16 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.device = device self.dtype = self.model_config.dtype + self.mm_config = self.model_config.multimodal_config + self.is_ec_producer_only = ( + self.vllm_config.ec_transfer_config is not None + and self.vllm_config.ec_transfer_config.is_ec_producer + and not self.vllm_config.ec_transfer_config.is_ec_consumer + ) + self.is_encoder_only = bool( + (self.mm_config is not None and self.mm_config.mm_encoder_only) + or self.is_ec_producer_only + ) self.kv_cache_dtype = self.dtype if self.cache_config.cache_dtype != "auto": # Quantized KV cache. @@ -237,6 +256,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) # KV Connector if configured. self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR + self.ec_connector = ECConnectorModelRunnerMixin() # For transferring state from execute_model to subsequent sample_tokens call. self.execute_model_state: ExecuteModelState | None = None @@ -332,6 +352,8 @@ def main_stream(self) -> torch.cuda.Stream: return torch.cuda.current_stream(self.device) def get_kv_cache_spec(self): + if self.is_encoder_only: + return {} return get_kv_cache_spec(self.vllm_config) def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: @@ -397,6 +419,10 @@ def _dummy_run( is_profile: bool = False, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if self.is_encoder_only: + empty = torch.empty(0, device=self.device) + return empty, empty + # Create a dummy scheduler output. num_reqs = min(num_tokens, self.max_num_reqs) if uniform_decode: @@ -509,6 +535,9 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: @torch.inference_mode() def profile_run(self) -> None: + if self.is_encoder_only: + return + hidden_states, sample_hidden_states = self._dummy_run( self.max_num_tokens, skip_attn=True, is_profile=True ) @@ -543,6 +572,9 @@ def profile_cudagraph_memory(self) -> int: @torch.inference_mode() def capture_model(self) -> int: + if self.is_encoder_only: + return 0 + if not self.cudagraph_manager.needs_capture(): logger.warning( "Skipping CUDA graph capture. To turn on CUDA graph capture, " @@ -920,9 +952,10 @@ def execute_model( self.update_requests(scheduler_output) self.block_tables.apply_staged_writes() if scheduler_output.total_num_scheduled_tokens == 0: + if self.is_encoder_only: + return make_empty_encoder_model_runner_output(scheduler_output) # No need to run the model. - empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self.kv_connector.no_forward(scheduler_output) # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -960,15 +993,17 @@ def execute_model( ) if batch_desc.num_tokens == 0: + if self.is_encoder_only: + return make_empty_encoder_model_runner_output(scheduler_output) # All DP ranks have zero tokens to run. - empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self.kv_connector.no_forward(scheduler_output) + block_tables = None + slot_mappings = None if not dummy_run: # Common case. # Prepare all the inputs and copy to the input buffers. input_batch = self.prepare_inputs(scheduler_output, batch_desc) - block_tables, slot_mappings = self.prepare_attn(input_batch) if self.lora_config: # Activate LoRA adapters. @@ -992,6 +1027,50 @@ def execute_model( slot_mappings = None # FIXME(woosuk): Fix warmup for LoRA. + inputs_embeds = None + if self.supports_mm_inputs and self.is_first_pp_rank: + assert self.encoder_cache is not None + encoder_outputs = self.encoder_cache.encoder_outputs + ec_context = nullcontext() + mm_hashes_to_save: list[str] = [] + if ( + not self.is_encoder_decoder + and scheduler_output.ec_connector_metadata is not None + ): + ec_context = self.ec_connector.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=encoder_outputs, + ) + if self.is_ec_producer_only: + mm_hashes_to_save, _ = ( + self.model_state.encoder_runner.prepare_mm_inputs( + scheduler_output.scheduled_encoder_inputs + ) + ) + + with ec_context as ec_connector_output: + # run MM encoder (if needed) and get multimodal embeddings. + inputs_embeds = self.model_state.get_mm_embeddings( + scheduler_output.scheduled_encoder_inputs, + input_batch, + self.req_states, + ) + for mm_hash in mm_hashes_to_save: + self.ec_connector.maybe_save_ec_to_connector( + encoder_outputs, + mm_hash, + ) + + if self.is_encoder_only: + output = make_empty_encoder_model_runner_output(scheduler_output) + output.ec_connector_output = ec_connector_output + return output + elif self.is_encoder_only: + return make_empty_encoder_model_runner_output(scheduler_output) + + if not dummy_run: + block_tables, slot_mappings = self.prepare_attn(input_batch) + attn_metadata = None slot_mappings_by_layer = None if not (dummy_run and skip_attn_for_dummy_run): @@ -1009,18 +1088,6 @@ def execute_model( self.kv_cache_config, ) - inputs_embeds = None - if self.supports_mm_inputs and self.is_first_pp_rank: - # Run MM encoder (if needed) and get multimodal embeddings. - # Only first PP rank prepares multimodal embeddings. - # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs - # to obtain inputs_embeds, because the compiled model expects this input. - inputs_embeds = self.model_state.get_mm_embeddings( - scheduler_output.scheduled_encoder_inputs, - input_batch, - self.req_states, - ) - model_inputs = { "input_ids": input_batch.input_ids, "positions": input_batch.positions, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 026b6a7d7eb9..9fec6ef7c7ab 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -33,6 +33,9 @@ def warmup_kernels( tokens each. The second iteration simulates a decode step with all requests generating 1 token each. """ + if model_runner.is_encoder_only: + return + prompt_token_ids = [0, 1] prompt_len = len(prompt_token_ids) num_spec_steps = model_runner.num_speculative_steps From d5ad859c965c3321e050ead9a0170a920c0e39d1 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Fri, 27 Mar 2026 19:23:59 +0000 Subject: [PATCH 02/31] update Signed-off-by: yewentao256 --- vllm/v1/worker/gpu/model_runner.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7de08c075b7a..a9ed726e82a2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1049,7 +1049,11 @@ def execute_model( ) with ec_context as ec_connector_output: - # run MM encoder (if needed) and get multimodal embeddings. + # Run MM encoder (if needed) and get multimodal embeddings. + # Only first PP rank prepares multimodal embeddings. + # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs + # to obtain inputs_embeds, because the compiled model + # expects this input. inputs_embeds = self.model_state.get_mm_embeddings( scheduler_output.scheduled_encoder_inputs, input_batch, From e434e5dd7d49e9e7c254b60977644eb49feb5196 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Fri, 27 Mar 2026 19:43:59 +0000 Subject: [PATCH 03/31] fix precommit Signed-off-by: yewentao256 --- vllm/v1/worker/gpu/model_runner.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a9ed726e82a2..892b8731796f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -20,7 +20,7 @@ import functools import gc import time -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from copy import deepcopy from typing import Any, NamedTuple @@ -47,6 +47,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.outputs import ( DraftTokenIds, + ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, make_empty_encoder_model_runner_output, @@ -1031,7 +1032,7 @@ def execute_model( if self.supports_mm_inputs and self.is_first_pp_rank: assert self.encoder_cache is not None encoder_outputs = self.encoder_cache.encoder_outputs - ec_context = nullcontext() + ec_context: AbstractContextManager[ECConnectorOutput | None] = nullcontext() # type: ignore[assignment] mm_hashes_to_save: list[str] = [] if ( not self.is_encoder_decoder From d7e63e8f69c4b23362da16624ce6d4f66d0deb5a Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Wed, 22 Jul 2026 21:10:24 +0000 Subject: [PATCH 04/31] update Signed-off-by: yewentao256 --- .../disaggregated_encoder/disagg_1e1pd_example.sh | 2 +- vllm/config/vllm.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh index ed752a38c6fd..4bf60ea66d61 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh @@ -29,7 +29,7 @@ if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then fi EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}" -TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}" # wait_for_server timeout +TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-300}" # wait_for_server timeout NUM_PROMPTS="${NUM_PROMPTS:-100}" # number of prompts to send in benchmark diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bda5013a5ff6..156d203cddda 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2207,10 +2207,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # Will be added by https://github.com/vllm-project/vllm/pull/35045 unsupported.append("KV sharing fast prefill") - if self.ec_transfer_config is not None: - # Will be added by https://github.com/vllm-project/vllm/pull/38390 - unsupported.append("EC transfer") - return unsupported def _validate_v2_model_runner(self) -> None: From 43d903d9320f378290d6b64b97aa093ac4372769 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Thu, 23 Jul 2026 15:25:34 +0300 Subject: [PATCH 05/31] Added Build Conector Worker Meta for EC Connector Signed-off-by: omerpaz95 --- .../ec_transfer/ec_connector/base.py | 33 +++++++++++++++++++ vllm/v1/outputs.py | 3 ++ .../worker/ec_connector_model_runner_mixin.py | 1 + 3 files changed, 37 insertions(+) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 1d5f467027e7..35b829a9f04f 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -20,6 +20,8 @@ get_finished() - called with ids of finished requests, returns ids of requests that have completed async sending/recving. + build_connector_worker_meta() - builds metadata to be sent + back to the scheduler-side connector """ import enum @@ -56,6 +58,27 @@ class ECConnectorMetadata(ABC): # noqa: B024 pass +class ECConnectorWorkerMetadata(ABC): + """ + Abstract Metadata used to communicate back + Worker ECConnector -> Scheduler ECConnector. + + Each worker can output its own metadata. + For a single engine step, all metadata objects returned by workers + will be aggregated using the `aggregate` method below, before + being passed to the Scheduler ECConnector. + """ + + @abstractmethod + def aggregate( + self, other: "ECConnectorWorkerMetadata" + ) -> "ECConnectorWorkerMetadata": + """ + Aggregate metadata with another `ECConnectorWorkerMetadata` object. + """ + pass + + class ECConnectorBase(ABC): def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): self._connector_metadata: ECConnectorMetadata | None = None @@ -190,6 +213,16 @@ def get_finished( """ return None, None + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: + """ + Build the ECConnector worker metadata for this engine step. + + Returns: + ECConnectorWorkerMetadata: the worker metadata. + None if no worker metadata is available. + """ + return None + # ============================== # Scheduler-side methods # ============================== diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 9f13ad939fc8..fa45fc2c548f 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -13,6 +13,7 @@ from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata from vllm.distributed.kv_events import KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorWorkerMetadata, @@ -22,6 +23,7 @@ KVConnectorStats = object KVConnectorWorkerMetadata = object KVConnectorKVEvents = object + ECConnectorWorkerMetadata = object class LogprobsLists(NamedTuple): @@ -226,6 +228,7 @@ class ECConnectorOutput: # [mm_hash] finished_sending: set[str] | None = None finished_recving: set[str] | None = None + ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None # ModelRunnerOutput is serialized and sent to the scheduler process. diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index b3430a8d94da..35054b2b8ff4 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -74,5 +74,6 @@ def _get_ec_connector_output( output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) ) + output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() From 251c90ee907c3e2d1e6ea78e2c012816bcf3f51b Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Thu, 23 Jul 2026 14:47:27 +0000 Subject: [PATCH 06/31] reduce mrv2 change Signed-off-by: yewentao256 --- .../worker/ec_connector_model_runner_mixin.py | 22 +++++++--- vllm/v1/worker/gpu/block_table.py | 2 + vllm/v1/worker/gpu/model_runner.py | 43 ++++++------------- 3 files changed, 31 insertions(+), 36 deletions(-) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index b3430a8d94da..899dd6cc512c 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -38,14 +38,21 @@ def maybe_save_ec_to_connector( def maybe_get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], + enabled: bool = True, + mm_hashes_to_save: list[str] | None = None, **kwargs, ) -> AbstractContextManager[ECConnectorOutput | None]: - return ( - ECConnectorModelRunnerMixin._get_ec_connector_output( - scheduler_output, encoder_cache, **kwargs - ) - if has_ec_transfer() - else nullcontext() + if ( + not enabled + or scheduler_output.ec_connector_metadata is None + or not has_ec_transfer() + ): + return nullcontext() + return ECConnectorModelRunnerMixin._get_ec_connector_output( + scheduler_output, + encoder_cache, + mm_hashes_to_save=mm_hashes_to_save, + **kwargs, ) # This context manager must be used within an active forward context. @@ -55,6 +62,7 @@ def maybe_get_ec_connector_output( def _get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], + mm_hashes_to_save: list[str] | None = None, **kwargs, ) -> Generator[ECConnectorOutput, None, None]: output = ECConnectorOutput() @@ -70,6 +78,8 @@ def _get_ec_connector_output( try: yield output + for mm_hash in mm_hashes_to_save or (): + ec_connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash) finally: output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 22c4afc11bc1..a09d3d24c82b 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -120,6 +120,8 @@ def append_block_ids( self.num_blocks.np[i, req_index] = start + len(block_ids) def apply_staged_writes(self) -> None: + if self.num_kv_cache_groups == 0: + return if self.num_kv_cache_groups == 1: # Single group: write directly, skipping the per-write group lookup. self.block_tables[0].apply_write() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 81d20e82bd6b..b94e03efed04 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -20,7 +20,6 @@ import functools import gc import time -from contextlib import AbstractContextManager, nullcontext from copy import deepcopy from typing import Any, NamedTuple @@ -57,7 +56,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import ( DraftTokenIds, - ECConnectorOutput, ModelRunnerOutput, make_empty_encoder_model_runner_output, ) @@ -1329,7 +1327,7 @@ def execute_model( assert self.encoder_cache is not None # Run MM encoder (if needed) and get multimodal embeddings. # Only first PP rank prepares multimodal embeddings. - ec_connector_output: ECConnectorOutput | None = None + ec_connector_output = None if dummy_run: # Obtain mm embeddings of correct shape for compiled model. inputs_embeds = self.model_state.dummy_inputs_embeds( @@ -1347,37 +1345,22 @@ def execute_model( scheduled_encoder_inputs=scheduled_encoder_inputs, ) - encoder_outputs = self.encoder_cache.encoder_outputs - ec_context: AbstractContextManager[ECConnectorOutput | None] = ( - nullcontext() - ) mm_hashes_to_save: list[str] = [] - if ( - not self.is_encoder_decoder - and scheduler_output.ec_connector_metadata is not None - ): - ec_context = self.ec_connector.maybe_get_ec_connector_output( - scheduler_output, - encoder_cache=encoder_outputs, - ) - if self.is_ec_producer_only: - mm_hashes_to_save, _ = ( - self.model_state.encoder_runner.prepare_mm_inputs( - scheduled_encoder_inputs - ) + if self.is_ec_producer_only: + mm_hashes_to_save, _ = ( + self.model_state.encoder_runner.prepare_mm_inputs( + scheduled_encoder_inputs ) - - with ec_context as ec_connector_output: + ) + with self.ec_connector.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache.encoder_outputs, + enabled=not self.is_encoder_decoder, + mm_hashes_to_save=mm_hashes_to_save, + ) as ec_connector_output: inputs_embeds = self.model_state.get_mm_embeddings( - scheduled_encoder_inputs, - input_batch, - self.req_states, + scheduled_encoder_inputs, input_batch, self.req_states ) - for mm_hash in mm_hashes_to_save: - self.ec_connector.maybe_save_ec_to_connector( - encoder_outputs, - mm_hash, - ) if self.is_encoder_only: output = make_empty_encoder_model_runner_output(scheduler_output) From 69b33fe11eef2124f34aa3fa1c9e4bcd148341ed Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Thu, 23 Jul 2026 15:04:11 +0000 Subject: [PATCH 07/31] reduce mrv2 code change Signed-off-by: yewentao256 --- .../worker/ec_connector_model_runner_mixin.py | 14 ++-- vllm/v1/worker/gpu/block_table.py | 4 + vllm/v1/worker/gpu/model_runner.py | 78 ++++++++----------- 3 files changed, 45 insertions(+), 51 deletions(-) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index 899dd6cc512c..fe4697ce86ea 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -39,7 +39,7 @@ def maybe_get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], enabled: bool = True, - mm_hashes_to_save: list[str] | None = None, + save_new_caches: bool = False, **kwargs, ) -> AbstractContextManager[ECConnectorOutput | None]: if ( @@ -51,7 +51,7 @@ def maybe_get_ec_connector_output( return ECConnectorModelRunnerMixin._get_ec_connector_output( scheduler_output, encoder_cache, - mm_hashes_to_save=mm_hashes_to_save, + save_new_caches=save_new_caches, **kwargs, ) @@ -62,7 +62,7 @@ def maybe_get_ec_connector_output( def _get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], - mm_hashes_to_save: list[str] | None = None, + save_new_caches: bool = False, **kwargs, ) -> Generator[ECConnectorOutput, None, None]: output = ECConnectorOutput() @@ -76,10 +76,14 @@ def _get_ec_connector_output( if ec_connector.is_consumer: ec_connector.start_load_caches(encoder_cache, **kwargs) + cached_hashes = set(encoder_cache) if save_new_caches else None try: yield output - for mm_hash in mm_hashes_to_save or (): - ec_connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash) + if cached_hashes is not None: + for mm_hash in encoder_cache.keys() - cached_hashes: + ec_connector.save_caches( + encoder_cache=encoder_cache, mm_hash=mm_hash + ) finally: output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index a09d3d24c82b..41855080f776 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -140,6 +140,8 @@ def gather_block_tables( out: tuple[torch.Tensor, ...] | None = None, out_ptrs: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...]: + if self.num_kv_cache_groups == 0: + return () if out is None: out = tuple(self.input_block_tables) out_ptrs = self.input_block_table_ptrs @@ -175,6 +177,8 @@ def compute_slot_mappings( num_tokens_padded: int, out: torch.Tensor | None = None, ) -> torch.Tensor: + if self.num_kv_cache_groups == 0: + return (self.slot_mappings if out is None else out)[:, :num_tokens_padded] num_reqs = idx_mapping.shape[0] num_groups = self.num_kv_cache_groups slot_mappings = self.slot_mappings if out is None else out diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b94e03efed04..5a4980b1fdf2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -129,7 +129,7 @@ logger = init_logger(__name__) -class GPUModelRunner(LoRAModelRunnerMixin): +class GPUModelRunner(LoRAModelRunnerMixin, ECConnectorModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -144,15 +144,15 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.device = device self.dtype = self.model_config.dtype - self.mm_config = self.model_config.multimodal_config + ec_config = vllm_config.ec_transfer_config self.is_ec_producer_only = ( - self.vllm_config.ec_transfer_config is not None - and self.vllm_config.ec_transfer_config.is_ec_producer - and not self.vllm_config.ec_transfer_config.is_ec_consumer + ec_config is not None + and ec_config.is_ec_producer + and not ec_config.is_ec_consumer ) - self.is_encoder_only = bool( - (self.mm_config is not None and self.mm_config.mm_encoder_only) - or self.is_ec_producer_only + mm_config = self.model_config.multimodal_config + self.is_encoder_only = self.is_ec_producer_only or bool( + mm_config and mm_config.mm_encoder_only ) self.kv_cache_dtype = self.dtype if self.cache_config.cache_dtype != "auto": @@ -272,7 +272,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # KV Connector if configured. self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR - self.ec_connector = ECConnectorModelRunnerMixin() # For transferring state from execute_model to subsequent sample_tokens call. self.execute_model_state: ExecuteModelState | None = None @@ -1218,10 +1217,9 @@ def execute_model( self.update_requests(scheduler_output) self.block_tables.apply_staged_writes() if scheduler_output.total_num_scheduled_tokens == 0: - if self.is_encoder_only: - return make_empty_encoder_model_runner_output(scheduler_output) # No need to run the model. - return self.kv_connector.no_forward(scheduler_output) + empty_output = self.kv_connector.no_forward(scheduler_output) + return empty_output # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -1255,29 +1253,25 @@ def execute_model( ) if batch_desc.num_tokens == 0: - if self.is_encoder_only: - return make_empty_encoder_model_runner_output(scheduler_output) # All DP ranks have zero tokens to run. - return self.kv_connector.no_forward(scheduler_output) + empty_output = self.kv_connector.no_forward(scheduler_output) + return empty_output - block_tables = None - slot_mappings = None if not dummy_run: # Common case. # Prepare all the inputs and copy to the input buffers. input_batch = self.prepare_inputs(scheduler_output, batch_desc) - if not self.is_encoder_only: - block_tables, slot_mappings = self.prepare_attn(input_batch) - # Mamba "align" pre-copy: migrate recurrent state across block - # boundaries before the forward. Runs only on real batches, and - # before model_state.prepare_attn gathers num_accepted_tokens so - # the boundary reset is visible to the attention metadata. - self.model_state.preprocess_state( - input_batch, - block_tables, - self.kv_cache_config, - self.req_states.num_computed_tokens.gpu, - ) + block_tables, slot_mappings = self.prepare_attn(input_batch) + # Mamba "align" pre-copy: migrate recurrent state across block + # boundaries before the forward. Runs only on real batches, and + # before model_state.prepare_attn gathers num_accepted_tokens so the + # boundary reset is visible to the attention metadata. + self.model_state.preprocess_state( + input_batch, + block_tables, + self.kv_cache_config, + self.req_states.num_computed_tokens.gpu, + ) if self.lora_config: # Activate LoRA adapters. @@ -1306,7 +1300,7 @@ def execute_model( attn_metadata = None slot_mappings_by_layer = None - if not self.is_encoder_only and not (dummy_run and skip_attn_for_dummy_run): + if not (dummy_run and skip_attn_for_dummy_run): assert slot_mappings is not None slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, self.kv_cache_config @@ -1323,11 +1317,11 @@ def execute_model( input_ids = input_batch.input_ids inputs_embeds = None + ec_connector_output = None if self.supports_mm_inputs and self.is_first_pp_rank: assert self.encoder_cache is not None # Run MM encoder (if needed) and get multimodal embeddings. # Only first PP rank prepares multimodal embeddings. - ec_connector_output = None if dummy_run: # Obtain mm embeddings of correct shape for compiled model. inputs_embeds = self.model_state.dummy_inputs_embeds( @@ -1345,31 +1339,23 @@ def execute_model( scheduled_encoder_inputs=scheduled_encoder_inputs, ) - mm_hashes_to_save: list[str] = [] - if self.is_ec_producer_only: - mm_hashes_to_save, _ = ( - self.model_state.encoder_runner.prepare_mm_inputs( - scheduled_encoder_inputs - ) - ) - with self.ec_connector.maybe_get_ec_connector_output( + with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache.encoder_outputs, enabled=not self.is_encoder_decoder, - mm_hashes_to_save=mm_hashes_to_save, + save_new_caches=self.is_ec_producer_only, ) as ec_connector_output: inputs_embeds = self.model_state.get_mm_embeddings( scheduled_encoder_inputs, input_batch, self.req_states ) - if self.is_encoder_only: - output = make_empty_encoder_model_runner_output(scheduler_output) - output.ec_connector_output = ec_connector_output - return output if inputs_embeds is not None and not self.model.requires_raw_input_tokens: input_ids = None - elif self.is_encoder_only: - return make_empty_encoder_model_runner_output(scheduler_output) + + if self.is_encoder_only: + output = make_empty_encoder_model_runner_output(scheduler_output) + output.ec_connector_output = ec_connector_output + return output model_inputs = { "input_ids": input_ids, From c3cf0f514f9b0b8bccc596d7924bd4b1224af34c Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Tue, 28 Jul 2026 16:20:29 +0300 Subject: [PATCH 08/31] Added a path for the empty step. Signed-off-by: omerpaz95 --- vllm/v1/outputs.py | 20 +++++++++++++++++++ .../worker/ec_connector_model_runner_mixin.py | 17 +++++++++++++++- vllm/v1/worker/gpu_model_runner.py | 16 +++++++++++++-- 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index e5f68b4d51a2..f840a5f83946 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -257,6 +257,13 @@ class ECConnectorOutput: finished_recving: set[str] | None = None ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None + def is_empty(self): + return ( + not self.finished_sending + and not self.finished_recving + and not self.ec_connector_worker_meta + ) + # ModelRunnerOutput is serialized and sent to the scheduler process. # This is expensive for torch.Tensor so prefer to use list instead. @@ -323,6 +330,19 @@ def with_kv_conn_output_only( output.kv_connector_output = kv_connector_output return output + @staticmethod + def with_ec_conn_output_only( + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return ModelRunnerOutput containing the provided ECConnectorOutput, + otherwise empty. + """ + if ec_connector_output is None or ec_connector_output.is_empty(): + return EMPTY_MODEL_RUNNER_OUTPUT + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index 35054b2b8ff4..4dc5a6f8a914 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -13,9 +13,10 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase from vllm.logger import init_logger -from vllm.v1.outputs import ECConnectorOutput +from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.v1.core.sched.output import SchedulerOutput logger = init_logger(__name__) @@ -34,6 +35,20 @@ def maybe_save_ec_to_connector( connector = get_ec_transfer() connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash) + @staticmethod + def ec_connector_no_forward( + scheduler_output: "SchedulerOutput", + vllm_config: "VllmConfig", + encoder_cache: dict[str, torch.Tensor], + ) -> ModelRunnerOutput: + # EC send/recv even if no work to do. + with ECConnectorModelRunnerMixin._get_ec_connector_output( + scheduler_output, encoder_cache=encoder_cache + ) as ec_connector_output: + pass + + return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output) + @staticmethod def maybe_get_ec_connector_output( scheduler_output: "SchedulerOutput", diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index dd31ddab672b..a6dfe703d41e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4221,10 +4221,22 @@ def execute_model( # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. self._dummy_run(1) - if not has_kv_transfer_group(): + if not has_kv_transfer_group() and not has_ec_transfer(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT - return self.kv_connector_no_forward(scheduler_output, self.vllm_config) + if has_kv_transfer_group(): + output = self.kv_connector_no_forward( + scheduler_output, self.vllm_config + ) + if has_ec_transfer(): + ec_output = self.ec_connector_no_forward( + scheduler_output, self.vllm_config, self.encoder_cache + ) + output.ec_connector_output = ec_output.ec_connector_output + return output + return self.ec_connector_no_forward( + scheduler_output, self.vllm_config, self.encoder_cache + ) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( From 3e22b0b5011ab80dc1e4438cdb51a530aa4cd6a2 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Tue, 28 Jul 2026 18:32:58 +0300 Subject: [PATCH 09/31] Added ECOutputAggregator for multiproc support. Signed-off-by: omerpaz95 --- .../ec_transfer/ec_connector/utils.py | 73 +++++++++++++++++++ vllm/v1/engine/core.py | 2 + vllm/v1/executor/abstract.py | 9 +++ vllm/v1/executor/multiproc_executor.py | 28 +++++-- 4 files changed, 106 insertions(+), 6 deletions(-) create mode 100644 vllm/distributed/ec_transfer/ec_connector/utils.py diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py new file mode 100644 index 000000000000..5eef9768e964 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EC connector helper utilities.""" + +from typing import TYPE_CHECKING + +from vllm.logger import init_logger +from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput + +if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase + +logger = init_logger(__name__) + + +class ECOutputAggregator: + """Utility class to aggregate the EC connector output of all workers + into a single output corresponding to `output_rank` for the scheduler. + + Mirrors KVOutputAggregator's role for KV connectors: only one worker's + ModelRunnerOutput (`output_rank`, e.g. the last pipeline-parallel rank) + reaches the scheduler, but the EC connector's real work may happen on a + different rank (e.g. the first PP rank, where the multimodal encoder + runs) -- this merges that rank's ec_connector_output onto the selected + output before it's returned. + """ + + def __init__(self, world_size: int): + self._world_size = world_size + + @classmethod + def from_connector(cls, connector: "ECConnectorBase", world_size: int): + return cls(world_size) + + def aggregate( + self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 + ) -> ModelRunnerOutput | None: + if not outputs[output_rank]: + return None + + finished_sending = set[str]() + finished_recving = set[str]() + aggregated_ec_connector_worker_meta = None + for model_runner_output in outputs: + assert model_runner_output is not None + ec_output = model_runner_output.ec_connector_output + if not ec_output: + continue + + finished_sending |= ec_output.finished_sending or set() + finished_recving |= ec_output.finished_recving or set() + + # Aggregate ec_connector_worker_meta from all workers. + if aggregated_ec_connector_worker_meta is None: + # Use the first worker's ec_connector_worker_meta as accumulator. + aggregated_ec_connector_worker_meta = ec_output.ec_connector_worker_meta + elif ec_connector_worker_meta := ec_output.ec_connector_worker_meta: + aggregated_ec_connector_worker_meta = ( + aggregated_ec_connector_worker_meta.aggregate( + ec_connector_worker_meta + ) + ) + + # select output of the worker specified by output_rank + output = outputs[output_rank] + + assert output is not None + output.ec_connector_output = ECConnectorOutput( + finished_sending=finished_sending or None, + finished_recving=finished_recving or None, + ec_connector_worker_meta=aggregated_ec_connector_worker_meta, + ) + return output diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9917f810b5b3..a95cc9121db0 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -172,6 +172,8 @@ def __init__( ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore + if self.scheduler.ec_connector is not None: # type: ignore + self.model_executor.init_ec_output_aggregator(self.scheduler.ec_connector) # type: ignore mm_registry = MULTIMODAL_REGISTRY self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config( diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 4063844d469c..2c1049ea076b 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -9,6 +9,7 @@ import vllm.envs as envs from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, @@ -25,6 +26,7 @@ from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase logger = init_logger(__name__) @@ -110,6 +112,7 @@ def __init__( self.is_sleeping = False self.sleeping_tags: set[str] = set() self.kv_output_aggregator: KVOutputAggregator | None = None + self.ec_output_aggregator: ECOutputAggregator | None = None @abstractmethod def _init_executor(self) -> None: @@ -283,6 +286,12 @@ def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: connector, self.parallel_config.world_size ) + def init_ec_output_aggregator(self, connector: "ECConnectorBase") -> None: + """Init ECOutputAggregator""" + self.ec_output_aggregator = ECOutputAggregator.from_connector( + connector, self.parallel_config.world_size + ) + @cached_property # Avoid unnecessary RPC calls def supported_tasks(self) -> tuple[SupportedTask, ...]: output: list[tuple[SupportedTask, ...]] diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 0a3a6d1ec369..20d3312f8ed9 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -29,6 +29,7 @@ from vllm.config import VllmConfig from vllm.distributed import destroy_distributed_environment, destroy_model_parallel from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.parallel_state import ( get_dcp_group, @@ -328,6 +329,7 @@ def execute_model( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def sample_tokens( # type: ignore[override] @@ -340,6 +342,7 @@ def sample_tokens( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def execute_dummy_batch(self) -> None: @@ -360,9 +363,11 @@ def collective_rpc( # type: ignore[override] non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator | None = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: - """Returns single result if unique_reply_rank and/or kv_output_aggregator - is provided, otherwise list.""" + """Returns single result if unique_reply_rank and/or an output + aggregator (kv_output_aggregator/ec_output_aggregator) is provided, + otherwise list.""" assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -372,11 +377,22 @@ def collective_rpc( # type: ignore[override] deadline = None if timeout is None else time.monotonic() + timeout kwargs = kwargs or {} - if kv_output_aggregator is not None: + aggregators = [ + agg + for agg in (kv_output_aggregator, ec_output_aggregator) + if agg is not None + ] + aggregate: Callable[[Any], Any] + if aggregators: output_rank = None - aggregate: Callable[[Any], Any] = partial( - kv_output_aggregator.aggregate, output_rank=unique_reply_rank or 0 - ) + + def _aggregate(outputs: Any) -> Any: + result = None + for agg in aggregators: + result = agg.aggregate(outputs, output_rank=unique_reply_rank or 0) + return result + + aggregate = _aggregate else: output_rank = unique_reply_rank aggregate = lambda x: x From 397cacc4b00ef4c430d749648bd404b23f2a602a Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Wed, 29 Jul 2026 10:55:35 +0300 Subject: [PATCH 10/31] Corrected PP semantics in EC output aggregation. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu_model_runner.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a6dfe703d41e..639c7aec5aa3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4224,19 +4224,24 @@ def execute_model( if not has_kv_transfer_group() and not has_ec_transfer(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT - if has_kv_transfer_group(): - output = self.kv_connector_no_forward( - scheduler_output, self.vllm_config - ) - if has_ec_transfer(): - ec_output = self.ec_connector_no_forward( - scheduler_output, self.vllm_config, self.encoder_cache - ) - output.ec_connector_output = ec_output.ec_connector_output - return output - return self.ec_connector_no_forward( - scheduler_output, self.vllm_config, self.encoder_cache + output = ( + self.kv_connector_no_forward(scheduler_output, self.vllm_config) + if has_kv_transfer_group() + else EMPTY_MODEL_RUNNER_OUTPUT ) + # EC transfer only ever runs on the first PP rank (that's + # where the multimodal encoder lives, see the is_first_rank + # gate above in _preprocess); other ranks have nothing of + # their own to report and must not touch encoder_cache. + if has_ec_transfer() and get_pp_group().is_first_rank: + ec_output = self.ec_connector_no_forward( + scheduler_output, self.vllm_config, self.encoder_cache + ) + if output is EMPTY_MODEL_RUNNER_OUTPUT: + # Don't mutate the shared singleton in place. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_output.ec_connector_output + return output if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( From 40b30ace1dc52b88b4d7ae21be5626e846b77532 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Wed, 29 Jul 2026 11:41:31 +0300 Subject: [PATCH 11/31] More bugfixing. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu_model_runner.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 639c7aec5aa3..29628f3ef66a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -941,6 +941,7 @@ def __init__( # Ephemeral state transferred between execute_model() and sample_tokens(). self.execute_model_state: ExecuteModelState | None = None self.kv_connector_output: KVConnectorOutput | None = None + self.ec_connector_output: ECConnectorOutput | None = None self.mamba_state_idx: dict[str, int] = {} self._mamba_bufs: mamba_utils.MambaBuffers | None = None self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None @@ -4479,6 +4480,7 @@ def execute_model( # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) self.kv_connector_output = kv_connector_output + self.ec_connector_output = ec_connector_output return hidden_states if self.is_pooling_model: @@ -4565,12 +4567,21 @@ def sample_tokens( if self.execute_model_state is None: kv_connector_output = self.kv_connector_output self.kv_connector_output = None + ec_connector_output = self.ec_connector_output + self.ec_connector_output = None # receive sampled token ids from the last PP rank. if self.use_async_scheduling and not get_pp_group().is_last_rank: self._pp_receive_prev_sampled_token_ids_to_input_batch() - # In case of PP with kv transfer, we need to pass through the - # kv_connector_output - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + # In case of PP with kv/ec transfer, we need to pass through their + # outputs -- this rank never has a "real" ModelRunnerOutput of its + # own (see the is_last_rank early return in execute_model above). + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + if ec_connector_output is not None and not ec_connector_output.is_empty(): + if output is EMPTY_MODEL_RUNNER_OUTPUT: + # Don't mutate the shared singleton in place. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output # Unpack ephemeral state. ( From 5cb71cc436e0eea212c0ec850be71f139bb5614a Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Thu, 30 Jul 2026 16:08:21 +0300 Subject: [PATCH 12/31] Fix ec_connector_output dropped on producer-only encoder-only step. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu_model_runner.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 29628f3ef66a..5b3925ccc5b4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4207,7 +4207,22 @@ def execute_model( encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) - return make_empty_encoder_model_runner_output(scheduler_output) + # Read ec_connector_output only after the context manager's + # __exit__ has run: that's what populates + # ec_connector_worker_meta (build_connector_worker_meta() is + # called in its finally block). Returning from inside the + # `with` would exit before that assignment lands, silently + # dropping the worker's completion report. + output = make_empty_encoder_model_runner_output(scheduler_output) + if ( + ec_connector_output is not None + and not ec_connector_output.is_empty() + ): + if output is EMPTY_MODEL_RUNNER_OUTPUT: + # Don't mutate the shared singleton in place. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output if not num_scheduled_tokens: if ( From 7e1f6cf72d2dbec2bced0ef0d4e69accd9119f27 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 3 Aug 2026 14:11:56 +0300 Subject: [PATCH 13/31] Added MRv2 support for EC Connector. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu/model_runner.py | 52 +++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0e74023e389a..b996940fdd97 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -20,7 +20,7 @@ import functools import gc import time -from copy import deepcopy +from copy import copy, deepcopy from typing import Any, NamedTuple import numpy as np @@ -31,6 +31,7 @@ from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.ec_transfer import has_ec_transfer from vllm.distributed.parallel_state import ( get_dcp_group, get_pp_group, @@ -56,7 +57,9 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds, + ECConnectorOutput, ModelRunnerOutput, make_empty_encoder_model_runner_output, ) @@ -1207,6 +1210,32 @@ def postprocess_sampled( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu ) + def _merge_ec_connector_no_forward( + self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput + ) -> ModelRunnerOutput: + """Merge the EC connector's send/recv bookkeeping into `output` for a + step with no work to run. The EC connector only ever runs on the + first PP rank (that's where the multimodal encoder lives); other + ranks have nothing of their own to report and must not touch + encoder_cache. + """ + if not ( + has_ec_transfer() + and self.is_first_pp_rank + and self.encoder_cache is not None + ): + return output + ec_connector_output = self.ec_connector_no_forward( + scheduler_output, self.vllm_config, self.encoder_cache.encoder_outputs + ).ec_connector_output + if ec_connector_output is None or ec_connector_output.is_empty(): + return output + if output is EMPTY_MODEL_RUNNER_OUTPUT: + # Don't mutate the shared singleton in place. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output + @torch.inference_mode() def execute_model( self, @@ -1227,7 +1256,9 @@ def execute_model( if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward( + scheduler_output, empty_output + ) # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -1263,7 +1294,7 @@ def execute_model( if batch_desc.num_tokens == 0: # All DP ranks have zero tokens to run. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward(scheduler_output, empty_output) if not dummy_run: # Common case. @@ -1458,6 +1489,7 @@ def execute_model( hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, finished_req_ids=finished_req_ids, + ec_connector_output=ec_connector_output, ) if not self.is_last_pp_rank: @@ -1480,6 +1512,7 @@ def sample_tokens( hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output self.execute_model_state = None if not self.is_last_pp_rank: @@ -1498,7 +1531,16 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + # This rank never has a "real" ModelRunnerOutput of its own (see the + # is_last_pp_rank early return in execute_model above), but may have + # produced ec_connector_output on the first PP rank -- pass it through. + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + if ec_connector_output is not None and not ec_connector_output.is_empty(): + if output is EMPTY_MODEL_RUNNER_OUTPUT: + # Don't mutate the shared singleton in place. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output # Last rank: sample tokens hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( @@ -1607,6 +1649,7 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) model_runner_output.kv_connector_output = kv_connector_output + model_runner_output.ec_connector_output = ec_connector_output return async_output @@ -1723,6 +1766,7 @@ class ExecuteModelState(NamedTuple): hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None finished_req_ids: set[str] + ec_connector_output: ECConnectorOutput | None def sort_batch_req_ids( From 2477323790645b72eb117e951a1aba17cf814633 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 3 Aug 2026 15:24:45 +0300 Subject: [PATCH 14/31] Fix EC connector no-forward crashing when ec_connector_metadata is None. ec_connector_no_forward() called _get_ec_connector_output() directly, skipping the ec_connector_metadata is None guard that maybe_get_ec_connector_output() has. Steps with no EC connector work scheduled hit the bare assert in _get_ec_connector_output() and crashed the EngineCore. Signed-off-by: omerpaz95 --- vllm/v1/worker/ec_connector_model_runner_mixin.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index 170d49106bf9..73e557304955 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -41,6 +41,10 @@ def ec_connector_no_forward( vllm_config: "VllmConfig", encoder_cache: dict[str, torch.Tensor], ) -> ModelRunnerOutput: + if scheduler_output.ec_connector_metadata is None: + # Nothing for the EC connector to do this step. + return ModelRunnerOutput.with_ec_conn_output_only(None) + # EC send/recv even if no work to do. with ECConnectorModelRunnerMixin._get_ec_connector_output( scheduler_output, encoder_cache=encoder_cache From 437d027d901c2da3431e1748419d3dc3415f7c1e Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Wed, 5 Aug 2026 15:57:40 +0300 Subject: [PATCH 15/31] Fixed a bug in saving EC Caches. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu/ec_connector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index 3a0a08f7c052..9d213b292bed 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -33,9 +33,11 @@ def __init__( encoder_cache: dict[str, torch.Tensor], ) -> None: self.encoder_cache = encoder_cache - self.save_new_caches = vllm_config.is_ec_producer_only self.ec_connector = get_ec_transfer() assert isinstance(self.ec_connector, ECConnectorBase) + # Every producer offloads freshly computed encoder outputs, including + # an ec_both node that also reloads them. + self.save_new_caches = self.ec_connector.is_producer @contextmanager def maybe_get_output( From cd54fad42b0de5bc0b7009903fda2cef84d09f8b Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Sun, 9 Aug 2026 11:06:12 +0300 Subject: [PATCH 16/31] Conformed to upstream model runner and fixed a bug in the return of empty model-runner output. Signed-off-by: omerpaz95 --- vllm/v1/outputs.py | 5 ++++- vllm/v1/worker/gpu/ec_connector.py | 26 +++++++++++++++++++++++++- vllm/v1/worker/gpu/model_runner.py | 4 ++-- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index f28f85862d61..6b1cecdbeb7b 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -373,7 +373,10 @@ def make_empty_encoder_model_runner_output( per-request bookkeeping but no generated data yet. """ if not scheduler_output.num_scheduled_tokens: - return EMPTY_MODEL_RUNNER_OUTPUT + # We don't want conusumers of this output to mutate shared + # module-level empty output object. + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + return output # Convert to list so we get a deterministic, indexable sequence req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys()) diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index 9d213b292bed..ffc5c077d8ed 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -9,7 +9,11 @@ from vllm.config import VllmConfig from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase -from vllm.v1.outputs import ECConnectorOutput +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + ModelRunnerOutput, +) if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -25,6 +29,12 @@ def maybe_get_output( ) -> Generator[ECConnectorOutput | None, None, None]: yield None + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + return EMPTY_MODEL_RUNNER_OUTPUT + class ActiveECConnector(ECConnector): def __init__( @@ -70,6 +80,20 @@ def maybe_get_output( output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + if scheduler_output.ec_connector_metadata is None: + # Nothing for the EC connector to do this step. + return ModelRunnerOutput.with_ec_conn_output_only(None) + + # EC send/recv even if no work to do. + with self.maybe_get_output(scheduler_output) as ec_connector_output: + pass + + return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output) + NO_OP_EC_CONNECTOR = ECConnector() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4c5390301390..0851e3436179 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1262,8 +1262,8 @@ def _merge_ec_connector_no_forward( and self.encoder_cache is not None ): return output - ec_connector_output = self.ec_connector_no_forward( - scheduler_output, self.vllm_config, self.encoder_cache.encoder_outputs + ec_connector_output = self.ec_connector.no_forward( + scheduler_output ).ec_connector_output if ec_connector_output is None or ec_connector_output.is_empty(): return output From ba2bc5a07a6a3b32eb438fbcb5979aaa9742b8ef Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Sun, 9 Aug 2026 17:16:21 +0300 Subject: [PATCH 17/31] Fixes ec_connector_output being dropped on pool() and on MRv2 encoder-only exit, which also mutated the shared EMPTY_MODEL_RUNNER_OUTPUT. Also added new helper and made the code prettier. Signed-off-by: omerpaz95 --- vllm/v1/outputs.py | 23 ++++++++++-- vllm/v1/worker/gpu/model_runner.py | 59 +++++++++++------------------- vllm/v1/worker/gpu_model_runner.py | 40 +++++++------------- 3 files changed, 55 insertions(+), 67 deletions(-) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 6b1cecdbeb7b..57a80b41f702 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -343,6 +343,24 @@ def with_ec_conn_output_only( output.ec_connector_output = ec_connector_output return output + @staticmethod + def attach_ec_conn_output( + output: "ModelRunnerOutput", + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return `output` carrying `ec_connector_output`. + + Sets the field on `output` and returns it, except that the shared + empty output is copied first rather than written to. Callers must use + the return value. + """ + if ec_connector_output is None or ec_connector_output.is_empty(): + return output + if output is EMPTY_MODEL_RUNNER_OUTPUT: + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): @@ -373,10 +391,7 @@ def make_empty_encoder_model_runner_output( per-request bookkeeping but no generated data yet. """ if not scheduler_output.num_scheduled_tokens: - # We don't want conusumers of this output to mutate shared - # module-level empty output object. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - return output + return EMPTY_MODEL_RUNNER_OUTPUT # Convert to list so we get a deterministic, indexable sequence req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys()) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0851e3436179..be14c7df6f87 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -20,7 +20,7 @@ import functools import gc import time -from copy import copy, deepcopy +from copy import deepcopy from typing import Any, NamedTuple import numpy as np @@ -31,7 +31,6 @@ from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.distributed.ec_transfer import has_ec_transfer from vllm.distributed.parallel_state import ( get_dcp_group, get_pp_group, @@ -61,7 +60,6 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import ( - EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds, ECConnectorOutput, ModelRunnerOutput, @@ -70,9 +68,6 @@ ) from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.cp_utils import check_attention_cp_compatibility -from vllm.v1.worker.ec_connector_model_runner_mixin import ( - ECConnectorModelRunnerMixin, -) from vllm.v1.worker.gpu import pcp_manager as pcp from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput from vllm.v1.worker.gpu.attn_utils import ( @@ -140,7 +135,7 @@ logger = init_logger(__name__) -class GPUModelRunner(LoRAModelRunnerMixin, ECConnectorModelRunnerMixin): +class GPUModelRunner(LoRAModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -1251,27 +1246,16 @@ def _merge_ec_connector_no_forward( self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput ) -> ModelRunnerOutput: """Merge the EC connector's send/recv bookkeeping into `output` for a - step with no work to run. The EC connector only ever runs on the - first PP rank (that's where the multimodal encoder lives); other - ranks have nothing of their own to report and must not touch - encoder_cache. + step with no work to run. + + A no-op unless this rank runs the EC connector: the connector is the + no-op one unless an encoder cache exists, which is only built on the + first PP rank of a multimodal model. """ - if not ( - has_ec_transfer() - and self.is_first_pp_rank - and self.encoder_cache is not None - ): - return output - ec_connector_output = self.ec_connector.no_forward( - scheduler_output - ).ec_connector_output - if ec_connector_output is None or ec_connector_output.is_empty(): - return output - if output is EMPTY_MODEL_RUNNER_OUTPUT: - # Don't mutate the shared singleton in place. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.attach_ec_conn_output( + output, + self.ec_connector.no_forward(scheduler_output).ec_connector_output, + ) @torch.inference_mode() def execute_model( @@ -1428,9 +1412,10 @@ def execute_model( input_ids = None if self.is_encoder_only: - output = make_empty_encoder_model_runner_output(scheduler_output) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.attach_ec_conn_output( + make_empty_encoder_model_runner_output(scheduler_output), + ec_connector_output, + ) model_inputs = { "input_ids": input_ids, @@ -1578,12 +1563,7 @@ def sample_tokens( # is_last_pp_rank early return in execute_model above), but may have # produced ec_connector_output on the first PP rank -- pass it through. output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) - if ec_connector_output is not None and not ec_connector_output.is_empty(): - if output is EMPTY_MODEL_RUNNER_OUTPUT: - # Don't mutate the shared singleton in place. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.attach_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( @@ -1712,6 +1692,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: input_batch = self.execute_model_state.input_batch hidden_states = self.execute_model_state.hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output self.execute_model_state = None # Post-step KV connector related operations. @@ -1719,7 +1700,10 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: if not self.is_last_pp_rank: self.postprocess_num_computed_tokens(input_batch) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.attach_ec_conn_output( + ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output), + ec_connector_output, + ) assert self.pooling_runner is not None pooler_output, finished_mask = self.pooling_runner.pool( @@ -1731,6 +1715,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: req_ids=input_batch.req_ids, req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, ) async_output = AsyncPoolingOutput( model_runner_output=model_runner_output, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 806a347236b5..b27e841af00f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3463,6 +3463,7 @@ def _pool( num_scheduled_tokens: int, num_scheduled_tokens_np: np.ndarray, kv_connector_output: KVConnectorOutput | None, + ec_connector_output: ECConnectorOutput | None, ) -> ModelRunnerOutput | AsyncModelRunnerOutput: num_reqs = self.input_batch.num_reqs assert num_reqs == len(self.input_batch.pooling_params), ( @@ -3500,6 +3501,7 @@ def _pool( req_ids=self.input_batch.req_ids.copy(), req_id_to_index=self.input_batch.req_id_to_index.copy(), kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, ) if raw_pooler_output is None or not any(finished_mask): @@ -4210,22 +4212,13 @@ def execute_model( encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) - # Read ec_connector_output only after the context manager's - # __exit__ has run: that's what populates - # ec_connector_worker_meta (build_connector_worker_meta() is - # called in its finally block). Returning from inside the - # `with` would exit before that assignment lands, silently - # dropping the worker's completion report. - output = make_empty_encoder_model_runner_output(scheduler_output) - if ( - ec_connector_output is not None - and not ec_connector_output.is_empty() - ): - if output is EMPTY_MODEL_RUNNER_OUTPUT: - # Don't mutate the shared singleton in place. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_connector_output - return output + # attach_ec_conn_output tests is_empty(), so it must run after + # the context manager's finally block has populated + # ec_connector_worker_meta. + return ModelRunnerOutput.attach_ec_conn_output( + make_empty_encoder_model_runner_output(scheduler_output), + ec_connector_output, + ) if not num_scheduled_tokens: if ( @@ -4256,10 +4249,9 @@ def execute_model( ec_output = self.ec_connector_no_forward( scheduler_output, self.vllm_config, self.encoder_cache ) - if output is EMPTY_MODEL_RUNNER_OUTPUT: - # Don't mutate the shared singleton in place. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_output.ec_connector_output + output = ModelRunnerOutput.attach_ec_conn_output( + output, ec_output.ec_connector_output + ) return output if self.cache_config.kv_sharing_fast_prefill: @@ -4501,6 +4493,7 @@ def execute_model( num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output, + ec_connector_output, ) sample_hidden_states = hidden_states[logits_indices] @@ -4587,12 +4580,7 @@ def sample_tokens( # outputs -- this rank never has a "real" ModelRunnerOutput of its # own (see the is_last_rank early return in execute_model above). output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) - if ec_connector_output is not None and not ec_connector_output.is_empty(): - if output is EMPTY_MODEL_RUNNER_OUTPUT: - # Don't mutate the shared singleton in place. - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.attach_ec_conn_output(output, ec_connector_output) # Unpack ephemeral state. ( From 491780fbf057168eeb2c3c7696dc8c12e491f5a7 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Sun, 9 Aug 2026 17:43:25 +0300 Subject: [PATCH 18/31] Fixed the EC Output Aggregator. Signed-off-by: omerpaz95 --- .../ec_transfer/ec_connector/utils.py | 57 ++++++------------- vllm/v1/engine/core.py | 2 +- vllm/v1/executor/abstract.py | 7 +-- vllm/v1/executor/multiproc_executor.py | 3 + 4 files changed, 24 insertions(+), 45 deletions(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index 5eef9768e964..45d2b5caeb29 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -2,45 +2,33 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """EC connector helper utilities.""" -from typing import TYPE_CHECKING - -from vllm.logger import init_logger from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput -if TYPE_CHECKING: - from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase - -logger = init_logger(__name__) - class ECOutputAggregator: - """Utility class to aggregate the EC connector output of all workers - into a single output corresponding to `output_rank` for the scheduler. + """Merge the EC connector output of all workers into the one output that + reaches the scheduler. Mirrors KVOutputAggregator's role for KV connectors: only one worker's ModelRunnerOutput (`output_rank`, e.g. the last pipeline-parallel rank) - reaches the scheduler, but the EC connector's real work may happen on a - different rank (e.g. the first PP rank, where the multimodal encoder - runs) -- this merges that rank's ec_connector_output onto the selected - output before it's returned. - """ - - def __init__(self, world_size: int): - self._world_size = world_size + reaches the scheduler, while the EC connector runs on the first pipeline + rank, where the multimodal encoder lives. - @classmethod - def from_connector(cls, connector: "ECConnectorBase", world_size: int): - return cls(world_size) + Finished sending/recving ids are unioned, so a connector that reports an + id from more than one worker must tolerate the scheduler acting on the + first report. + """ def aggregate( self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 ) -> ModelRunnerOutput | None: - if not outputs[output_rank]: + output = outputs[output_rank] + if not output: return None finished_sending = set[str]() finished_recving = set[str]() - aggregated_ec_connector_worker_meta = None + worker_meta = None for model_runner_output in outputs: assert model_runner_output is not None ec_output = model_runner_output.ec_connector_output @@ -50,24 +38,15 @@ def aggregate( finished_sending |= ec_output.finished_sending or set() finished_recving |= ec_output.finished_recving or set() - # Aggregate ec_connector_worker_meta from all workers. - if aggregated_ec_connector_worker_meta is None: - # Use the first worker's ec_connector_worker_meta as accumulator. - aggregated_ec_connector_worker_meta = ec_output.ec_connector_worker_meta - elif ec_connector_worker_meta := ec_output.ec_connector_worker_meta: - aggregated_ec_connector_worker_meta = ( - aggregated_ec_connector_worker_meta.aggregate( - ec_connector_worker_meta - ) - ) - - # select output of the worker specified by output_rank - output = outputs[output_rank] + if worker_meta is None: + worker_meta = ec_output.ec_connector_worker_meta + elif other := ec_output.ec_connector_worker_meta: + worker_meta = worker_meta.aggregate(other) - assert output is not None - output.ec_connector_output = ECConnectorOutput( + aggregated = ECConnectorOutput( finished_sending=finished_sending or None, finished_recving=finished_recving or None, - ec_connector_worker_meta=aggregated_ec_connector_worker_meta, + ec_connector_worker_meta=worker_meta, ) + output.ec_connector_output = None if aggregated.is_empty() else aggregated return output diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ed520a12819f..25413dbd1b74 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -173,7 +173,7 @@ def __init__( if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore if self.scheduler.ec_connector is not None: # type: ignore - self.model_executor.init_ec_output_aggregator(self.scheduler.ec_connector) # type: ignore + self.model_executor.init_ec_output_aggregator() mm_registry = MULTIMODAL_REGISTRY self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config( diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index a6a48c021ce8..0b570e2f38a6 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -26,7 +26,6 @@ from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase if TYPE_CHECKING: - from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase logger = init_logger(__name__) @@ -286,11 +285,9 @@ def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: connector, self.parallel_config.world_size ) - def init_ec_output_aggregator(self, connector: "ECConnectorBase") -> None: + def init_ec_output_aggregator(self) -> None: """Init ECOutputAggregator""" - self.ec_output_aggregator = ECOutputAggregator.from_connector( - connector, self.parallel_config.world_size - ) + self.ec_output_aggregator = ECOutputAggregator() @cached_property # Avoid unnecessary RPC calls def supported_tasks(self) -> tuple[SupportedTask, ...]: diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 39b353cd9d48..00b60c2b43bd 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -398,6 +398,9 @@ def collective_rpc( # type: ignore[override] output_rank = None def _aggregate(outputs: Any) -> Any: + # Each aggregator merges its own connector's output onto + # outputs[output_rank] in place and returns it, so chaining + # them and keeping the last result is safe. result = None for agg in aggregators: result = agg.aggregate(outputs, output_rank=unique_reply_rank or 0) From 450f0311fc3cf90b7bd3c7015f5c9cfb9d6b07fe Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 12:11:09 +0300 Subject: [PATCH 19/31] Removed un-needed assert, fixed docstrings. Signed-off-by: omerpaz95 --- vllm/distributed/ec_transfer/ec_connector/utils.py | 14 ++++---------- vllm/v1/worker/gpu/model_runner.py | 6 ++---- 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index 45d2b5caeb29..e9cd57115dda 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -6,17 +6,11 @@ class ECOutputAggregator: - """Merge the EC connector output of all workers into the one output that - reaches the scheduler. + """Merge every worker's EC connector output into the one ModelRunnerOutput + that reaches the scheduler. - Mirrors KVOutputAggregator's role for KV connectors: only one worker's - ModelRunnerOutput (`output_rank`, e.g. the last pipeline-parallel rank) - reaches the scheduler, while the EC connector runs on the first pipeline - rank, where the multimodal encoder lives. - - Finished sending/recving ids are unioned, so a connector that reports an - id from more than one worker must tolerate the scheduler acting on the - first report. + Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the + scheduler, but the EC connector may have run on other ranks. """ def aggregate( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index be14c7df6f87..0b3f0622f8f6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1245,12 +1245,11 @@ def postprocess_sampled( def _merge_ec_connector_no_forward( self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput ) -> ModelRunnerOutput: - """Merge the EC connector's send/recv bookkeeping into `output` for a + """Merge the EC connector's output into `output` for a step with no work to run. A no-op unless this rank runs the EC connector: the connector is the - no-op one unless an encoder cache exists, which is only built on the - first PP rank of a multimodal model. + no-op one unless an encoder cache exists. """ return ModelRunnerOutput.attach_ec_conn_output( output, @@ -1383,7 +1382,6 @@ def execute_model( inputs_embeds = None ec_connector_output = None if self.supports_mm_inputs and self.is_first_pp_rank: - assert self.encoder_cache is not None # Run MM encoder (if needed) and get multimodal embeddings. # Only first PP rank prepares multimodal embeddings. if dummy_run: From fc139b6e24995b30a8a51b876385e6bb4ea10500 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 12:15:24 +0300 Subject: [PATCH 20/31] changed method name (attach_ec_conn_output -> with_ec_conn_output) Signed-off-by: omerpaz95 --- vllm/v1/outputs.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 8 ++++---- vllm/v1/worker/gpu_model_runner.py | 8 ++++---- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 57a80b41f702..564a80e41731 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -344,7 +344,7 @@ def with_ec_conn_output_only( return output @staticmethod - def attach_ec_conn_output( + def with_ec_conn_output( output: "ModelRunnerOutput", ec_connector_output: ECConnectorOutput | None, ) -> "ModelRunnerOutput": diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0b3f0622f8f6..270c4269017a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1251,7 +1251,7 @@ def _merge_ec_connector_no_forward( A no-op unless this rank runs the EC connector: the connector is the no-op one unless an encoder cache exists. """ - return ModelRunnerOutput.attach_ec_conn_output( + return ModelRunnerOutput.with_ec_conn_output( output, self.ec_connector.no_forward(scheduler_output).ec_connector_output, ) @@ -1410,7 +1410,7 @@ def execute_model( input_ids = None if self.is_encoder_only: - return ModelRunnerOutput.attach_ec_conn_output( + return ModelRunnerOutput.with_ec_conn_output( make_empty_encoder_model_runner_output(scheduler_output), ec_connector_output, ) @@ -1561,7 +1561,7 @@ def sample_tokens( # is_last_pp_rank early return in execute_model above), but may have # produced ec_connector_output on the first PP rank -- pass it through. output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) - return ModelRunnerOutput.attach_ec_conn_output(output, ec_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( @@ -1698,7 +1698,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: if not self.is_last_pp_rank: self.postprocess_num_computed_tokens(input_batch) - return ModelRunnerOutput.attach_ec_conn_output( + return ModelRunnerOutput.with_ec_conn_output( ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output), ec_connector_output, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b27e841af00f..bb39658cd5cd 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4212,10 +4212,10 @@ def execute_model( encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) - # attach_ec_conn_output tests is_empty(), so it must run after + # with_ec_conn_output tests is_empty(), so it must run after # the context manager's finally block has populated # ec_connector_worker_meta. - return ModelRunnerOutput.attach_ec_conn_output( + return ModelRunnerOutput.with_ec_conn_output( make_empty_encoder_model_runner_output(scheduler_output), ec_connector_output, ) @@ -4249,7 +4249,7 @@ def execute_model( ec_output = self.ec_connector_no_forward( scheduler_output, self.vllm_config, self.encoder_cache ) - output = ModelRunnerOutput.attach_ec_conn_output( + output = ModelRunnerOutput.with_ec_conn_output( output, ec_output.ec_connector_output ) return output @@ -4580,7 +4580,7 @@ def sample_tokens( # outputs -- this rank never has a "real" ModelRunnerOutput of its # own (see the is_last_rank early return in execute_model above). output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) - return ModelRunnerOutput.attach_ec_conn_output(output, ec_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Unpack ephemeral state. ( From df6daf0ad02f47f6f51cccf9f3db3bd89f94c773 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 12:46:15 +0300 Subject: [PATCH 21/31] Some nit fixes and removed redundant diffs. Signed-off-by: omerpaz95 --- .../worker/ec_connector_model_runner_mixin.py | 30 +++---------------- vllm/v1/worker/gpu_model_runner.py | 2 +- 2 files changed, 5 insertions(+), 27 deletions(-) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index 73e557304955..afe95c59b732 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -16,7 +16,6 @@ from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: - from vllm.config import VllmConfig from vllm.v1.core.sched.output import SchedulerOutput logger = init_logger(__name__) @@ -38,16 +37,11 @@ def maybe_save_ec_to_connector( @staticmethod def ec_connector_no_forward( scheduler_output: "SchedulerOutput", - vllm_config: "VllmConfig", encoder_cache: dict[str, torch.Tensor], ) -> ModelRunnerOutput: - if scheduler_output.ec_connector_metadata is None: - # Nothing for the EC connector to do this step. - return ModelRunnerOutput.with_ec_conn_output_only(None) - # EC send/recv even if no work to do. - with ECConnectorModelRunnerMixin._get_ec_connector_output( - scheduler_output, encoder_cache=encoder_cache + with ECConnectorModelRunnerMixin.maybe_get_ec_connector_output( + scheduler_output, encoder_cache ) as ec_connector_output: pass @@ -57,21 +51,12 @@ def ec_connector_no_forward( def maybe_get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], - enabled: bool = True, - save_new_caches: bool = False, **kwargs, ) -> AbstractContextManager[ECConnectorOutput | None]: - if ( - not enabled - or scheduler_output.ec_connector_metadata is None - or not has_ec_transfer() - ): + if scheduler_output.ec_connector_metadata is None or not has_ec_transfer(): return nullcontext() return ECConnectorModelRunnerMixin._get_ec_connector_output( - scheduler_output, - encoder_cache, - save_new_caches=save_new_caches, - **kwargs, + scheduler_output, encoder_cache, **kwargs ) # This context manager must be used within an active forward context. @@ -81,7 +66,6 @@ def maybe_get_ec_connector_output( def _get_ec_connector_output( scheduler_output: "SchedulerOutput", encoder_cache: dict[str, torch.Tensor], - save_new_caches: bool = False, **kwargs, ) -> Generator[ECConnectorOutput, None, None]: output = ECConnectorOutput() @@ -95,14 +79,8 @@ def _get_ec_connector_output( if ec_connector.is_consumer: ec_connector.start_load_caches(encoder_cache, **kwargs) - cached_hashes = set(encoder_cache) if save_new_caches else None try: yield output - if cached_hashes is not None: - for mm_hash in encoder_cache.keys() - cached_hashes: - ec_connector.save_caches( - encoder_cache=encoder_cache, mm_hash=mm_hash - ) finally: output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bb39658cd5cd..eda63a596061 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4247,7 +4247,7 @@ def execute_model( # their own to report and must not touch encoder_cache. if has_ec_transfer() and get_pp_group().is_first_rank: ec_output = self.ec_connector_no_forward( - scheduler_output, self.vllm_config, self.encoder_cache + scheduler_output, self.encoder_cache ) output = ModelRunnerOutput.with_ec_conn_output( output, ec_output.ec_connector_output From e7a2da75122f8b8a9de461ded891dba116feae45 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 12:58:57 +0300 Subject: [PATCH 22/31] Shortened diff. Signed-off-by: omerpaz95 --- vllm/v1/worker/ec_connector_model_runner_mixin.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index afe95c59b732..bbdf41c3362f 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -53,10 +53,12 @@ def maybe_get_ec_connector_output( encoder_cache: dict[str, torch.Tensor], **kwargs, ) -> AbstractContextManager[ECConnectorOutput | None]: - if scheduler_output.ec_connector_metadata is None or not has_ec_transfer(): - return nullcontext() - return ECConnectorModelRunnerMixin._get_ec_connector_output( - scheduler_output, encoder_cache, **kwargs + return ( + ECConnectorModelRunnerMixin._get_ec_connector_output( + scheduler_output, encoder_cache, **kwargs + ) + if has_ec_transfer() and scheduler_output.ec_connector_metadata is not None + else nullcontext() ) # This context manager must be used within an active forward context. From 8d50582f257383f6a4f6a8d2129b0dda75bcf6df Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 14:52:23 +0300 Subject: [PATCH 23/31] Removed redundant check. Signed-off-by: omerpaz95 --- vllm/v1/worker/gpu/ec_connector.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index ffc5c077d8ed..5dc8d92359a9 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -84,10 +84,6 @@ def no_forward( self, scheduler_output: "SchedulerOutput", ) -> ModelRunnerOutput: - if scheduler_output.ec_connector_metadata is None: - # Nothing for the EC connector to do this step. - return ModelRunnerOutput.with_ec_conn_output_only(None) - # EC send/recv even if no work to do. with self.maybe_get_output(scheduler_output) as ec_connector_output: pass From dad0500e91412b8a5dbbc6b756c58d03fd3294df Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Mon, 10 Aug 2026 14:57:43 +0300 Subject: [PATCH 24/31] Added unit tests. Signed-off-by: omerpaz95 --- .buildkite/test_areas/misc.yaml | 1 + .../unit/test_ec_output_aggregator.py | 125 +++++++++++ .../unit/test_worker_ec_connector.py | 200 ++++++++++++++++++ tests/v1/test_outputs.py | 41 +++- 4 files changed, 366 insertions(+), 1 deletion(-) create mode 100644 tests/v1/ec_connector/unit/test_ec_output_aggregator.py create mode 100644 tests/v1/ec_connector/unit/test_worker_ec_connector.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cb5bc921aa98..c70f3ab2d66a 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -124,6 +124,7 @@ steps: - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/ec_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics - label: Extract Hidden States Integration diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py new file mode 100644 index 000000000000..d50647ca8840 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ECOutputAggregator.""" + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator +from vllm.v1.outputs import ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput + +pytestmark = pytest.mark.cpu_test + + +class FakeWorkerMeta(ECConnectorWorkerMetadata): + """Per-worker save/load reports, concatenated in merge order. + + `aggregate` returns a new object, as the base class declares: an + aggregator that discarded the return value would lose the merge. + """ + + def __init__( + self, + completed_saves: list[str] | None = None, + completed_loads: list[str] | None = None, + ): + self.completed_saves = completed_saves or [] + self.completed_loads = completed_loads or [] + + def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": + return FakeWorkerMeta( + self.completed_saves + other.completed_saves, + self.completed_loads + other.completed_loads, + ) + + +def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[], req_id_to_index={}, ec_connector_output=ec_output + ) + + +def test_aggregate_unions_finished_ids_onto_output_rank(): + """EC work done on any rank reaches the scheduler via output_rank's output.""" + outputs = [ + _worker_output( + ECConnectorOutput(finished_sending={"mm0"}, finished_recving={"mm1"}) + ), + _worker_output(ECConnectorOutput(finished_sending={"mm2"})), + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=1) + + assert result is outputs[1] + assert result.ec_connector_output.finished_sending == {"mm0", "mm2"} + assert result.ec_connector_output.finished_recving == {"mm1"} + + +def test_aggregate_worker_meta_folds_across_ranks(): + """Worker metadata is folded left across ranks, keeping each merge result.""" + outputs = [ + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"], [])) + ), + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta([], ["mm1"])) + ), + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"], ["mm3"])) + ), + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=2) + + worker_meta = result.ec_connector_output.ec_connector_worker_meta + assert worker_meta.completed_saves == ["mm0", "mm2"] + assert worker_meta.completed_loads == ["mm1", "mm3"] + + +def test_aggregate_worker_meta_tolerates_ranks_without_meta(): + """Ranks reporting no metadata neither seed nor clobber the accumulator.""" + worker_meta = FakeWorkerMeta(["mm1"], []) + outputs = [ + _worker_output(ECConnectorOutput(finished_sending={"mm0"})), + _worker_output(ECConnectorOutput(ec_connector_worker_meta=worker_meta)), + _worker_output(ECConnectorOutput(finished_recving={"mm2"})), + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=0) + + assert result.ec_connector_output.ec_connector_worker_meta is worker_meta + + +def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): + """Empty per-worker reports must not reach the scheduler as an empty object.""" + outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())] + + result = ECOutputAggregator().aggregate(outputs, output_rank=0) + + assert result is outputs[0] + assert result.ec_connector_output is None + + assert ECOutputAggregator().aggregate([None], output_rank=0) is None + + +def test_chaining_with_kv_aggregator_preserves_both_outputs(): + """MultiprocExecutor chains both aggregators and keeps only the last result, + so each must merge onto the same output_rank output rather than replace it. + """ + outputs = [ + _worker_output(ECConnectorOutput(finished_sending={"mm0"})), + _worker_output(None), + ] + outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7}) + + result = None + for aggregator in ( + KVOutputAggregator(expected_finished_count=1), + ECOutputAggregator(), + ): + result = aggregator.aggregate(outputs, output_rank=1) + + assert result is outputs[1] + assert result.kv_connector_output.invalid_block_ids == {7} + assert result.ec_connector_output.finished_sending == {"mm0"} diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py new file mode 100644 index 000000000000..3834bdf60c9a --- /dev/null +++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU model runner's EC connector wrapper.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, + ECConnectorWorkerMetadata, +) +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.v1.worker.gpu.ec_connector import ( + NO_OP_EC_CONNECTOR, + ActiveECConnector, + get_ec_connector, +) + +pytestmark = pytest.mark.cpu_test + + +class FakeMetadata(ECConnectorMetadata): + pass + + +class FakeWorkerMeta(ECConnectorWorkerMetadata): + """Mirrors ECCPUWorkerMetadata: the transfers that completed this step.""" + + def __init__( + self, + completed_saves: list[str] | None = None, + completed_loads: list[str] | None = None, + ): + self.completed_saves = completed_saves or [] + self.completed_loads = completed_loads or [] + + def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": + return self + + +class FakeECConnector(ECConnectorBase): + """Records the calls that ActiveECConnector is expected to drive.""" + + def __init__(self, is_producer: bool = True, is_consumer: bool = False): + # ECConnectorBase.__init__ requires a full VllmConfig; set only what + # the worker-side code reads. + self._connector_metadata = None + self._is_producer = is_producer + self._is_consumer = is_consumer + self.bound_metadata: list[ECConnectorMetadata] = [] + self.clear_calls = 0 + self.load_calls: list[dict] = [] + self.saved_hashes: list[str] = [] + self.worker_meta_calls = 0 + self.worker_meta: FakeWorkerMeta | None = None + self.finished: tuple[set[str] | None, set[str] | None] = (None, None) + + def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> None: + self.bound_metadata.append(connector_metadata) + super().bind_connector_metadata(connector_metadata) + + def clear_connector_metadata(self) -> None: + self.clear_calls += 1 + super().clear_connector_metadata() + + def start_load_caches(self, encoder_cache: dict, **kwargs) -> None: + self.load_calls.append(encoder_cache) + + def save_caches(self, encoder_cache: dict, mm_hash: str) -> None: + self.saved_hashes.append(mm_hash) + + def get_finished(self, finished_req_ids: set[str]): + return self.finished + + def build_connector_worker_meta(self) -> FakeWorkerMeta | None: + self.worker_meta_calls += 1 + return self.worker_meta + + # Scheduler-side abstract methods, unused by these tests. + def has_cache_item(self, identifier: str) -> bool: + return False + + def update_state_after_alloc(self, request, index: int) -> None: + pass + + def build_connector_meta(self, scheduler_output) -> ECConnectorMetadata: + return FakeMetadata() + + +def _scheduler_output(metadata: ECConnectorMetadata | None) -> SimpleNamespace: + return SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=frozenset()) + + +def _active_connector( + encoder_cache: dict | None = None, **connector_kwargs +) -> tuple[ActiveECConnector, FakeECConnector]: + fake = FakeECConnector(**connector_kwargs) + with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): + connector = ActiveECConnector(SimpleNamespace(), encoder_cache or {}) + return connector, fake + + +def test_no_forward_is_noop_without_ec_connector(): + """`ec_connector_metadata is None` means no EC connector is configured: + nothing to poll, and the shared empty output is returned as-is. + """ + scheduler_output = _scheduler_output(metadata=None) + + assert NO_OP_EC_CONNECTOR.no_forward(scheduler_output) is EMPTY_MODEL_RUNNER_OUTPUT + + connector, fake = _active_connector() + assert connector.no_forward(scheduler_output) is EMPTY_MODEL_RUNNER_OUTPUT + assert fake.bound_metadata == [] + assert fake.worker_meta_calls == 0 + + +def test_no_forward_polls_connector_with_empty_metadata(): + """A step with no work still reaps completed transfers. + + The scheduler sends metadata every step, empty or not, and that is what + drives build_connector_worker_meta() -- the only channel by which finished + saves and loads are reported. + """ + connector, fake = _active_connector() + fake.worker_meta = FakeWorkerMeta(completed_saves=["mm0"]) + scheduler_output = _scheduler_output(metadata=FakeMetadata()) + + output = connector.no_forward(scheduler_output) + + assert fake.worker_meta_calls == 1 + assert output.ec_connector_output.ec_connector_worker_meta is fake.worker_meta + assert fake.bound_metadata == [scheduler_output.ec_connector_metadata] + assert fake.clear_calls == 1 + + +@pytest.mark.parametrize( + ("is_producer", "is_consumer"), + [(True, False), (True, True), (False, True)], +) +def test_maybe_get_output_saves_only_newly_added_caches(is_producer, is_consumer): + """Every producer offloads caches computed during the step, including an + ec_both node, and never re-offloads ones that were already cached. + """ + encoder_cache = {"mm_old": None} + connector, fake = _active_connector( + encoder_cache, is_producer=is_producer, is_consumer=is_consumer + ) + + with connector.maybe_get_output(_scheduler_output(FakeMetadata())): + encoder_cache["mm_new"] = None + + assert fake.saved_hashes == (["mm_new"] if is_producer else []) + assert fake.load_calls == ([encoder_cache] if is_consumer else []) + + +def test_worker_meta_is_populated_only_on_exit(): + """The worker's report lands in the finally block, so a caller returning + from inside the `with` would drop it. + """ + connector, fake = _active_connector() + fake.worker_meta = FakeWorkerMeta(completed_loads=["mm0"]) + + with connector.maybe_get_output(_scheduler_output(FakeMetadata())) as output: + assert output.ec_connector_worker_meta is None + assert fake.worker_meta_calls == 0 + + assert output.ec_connector_worker_meta is fake.worker_meta + + +def test_get_ec_connector_activates_only_for_a_multimodal_ec_deployment(): + encoder_cache = SimpleNamespace(encoder_outputs={}) + config = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False)) + encoder_decoder = SimpleNamespace( + model_config=SimpleNamespace(is_encoder_decoder=True) + ) + + with patch( + "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", + return_value=FakeECConnector(), + ): + with patch( + "vllm.v1.worker.gpu.ec_connector.has_ec_transfer", return_value=False + ): + assert get_ec_connector(config, encoder_cache) is NO_OP_EC_CONNECTOR + + with patch( + "vllm.v1.worker.gpu.ec_connector.has_ec_transfer", return_value=True + ): + assert get_ec_connector(encoder_decoder, encoder_cache) is ( + NO_OP_EC_CONNECTOR + ) + assert get_ec_connector(config, None) is NO_OP_EC_CONNECTOR + connector = get_ec_connector(config, encoder_cache) + + assert isinstance(connector, ActiveECConnector) + # Aliases the live dict, so caches added later in the step are visible. + assert connector.encoder_cache is encoder_cache.encoder_outputs diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 4696eefa29fe..c9322fd29409 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -4,7 +4,13 @@ import torch -from vllm.v1.outputs import LogprobsLists, LogprobsTensors +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + LogprobsLists, + LogprobsTensors, + ModelRunnerOutput, +) def test_logprobs_tensors_cat(): @@ -122,3 +128,36 @@ def test_slice_all_requests(self): assert len(sliced.logprob_token_ids) == 9 # All tokens assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids assert sliced.cu_num_generated_tokens is None + + +def _ec_output() -> ECConnectorOutput: + return ECConnectorOutput(finished_sending={"mm_hash"}) + + +def test_with_ec_conn_output_keeps_output_when_nothing_to_report(): + output = ModelRunnerOutput(req_ids=["r0"], req_id_to_index={"r0": 0}) + + assert ModelRunnerOutput.with_ec_conn_output(output, None) is output + assert ModelRunnerOutput.with_ec_conn_output(output, ECConnectorOutput()) is output + assert output.ec_connector_output is None + + +def test_with_ec_conn_output_sets_field_in_place(): + output = ModelRunnerOutput(req_ids=["r0"], req_id_to_index={"r0": 0}) + ec_output = _ec_output() + + result = ModelRunnerOutput.with_ec_conn_output(output, ec_output) + + assert result is output + assert result.ec_connector_output is ec_output + + +def test_with_ec_conn_output_copies_shared_empty_output(): + """The shared empty output is copied, never written to.""" + ec_output = _ec_output() + + result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output) + + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output is ec_output + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None From 08d29c7552e9e3d421e95059b229c67087081128 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Tue, 11 Aug 2026 11:56:47 +0300 Subject: [PATCH 25/31] Reverted MRv1 changes. Signed-off-by: omerpaz95 --- .../worker/ec_connector_model_runner_mixin.py | 18 +------- vllm/v1/worker/gpu_model_runner.py | 43 +++---------------- 2 files changed, 8 insertions(+), 53 deletions(-) diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py index bbdf41c3362f..b3430a8d94da 100644 --- a/vllm/v1/worker/ec_connector_model_runner_mixin.py +++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py @@ -13,7 +13,7 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase from vllm.logger import init_logger -from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput +from vllm.v1.outputs import ECConnectorOutput if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -34,19 +34,6 @@ def maybe_save_ec_to_connector( connector = get_ec_transfer() connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash) - @staticmethod - def ec_connector_no_forward( - scheduler_output: "SchedulerOutput", - encoder_cache: dict[str, torch.Tensor], - ) -> ModelRunnerOutput: - # EC send/recv even if no work to do. - with ECConnectorModelRunnerMixin.maybe_get_ec_connector_output( - scheduler_output, encoder_cache - ) as ec_connector_output: - pass - - return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output) - @staticmethod def maybe_get_ec_connector_output( scheduler_output: "SchedulerOutput", @@ -57,7 +44,7 @@ def maybe_get_ec_connector_output( ECConnectorModelRunnerMixin._get_ec_connector_output( scheduler_output, encoder_cache, **kwargs ) - if has_ec_transfer() and scheduler_output.ec_connector_metadata is not None + if has_ec_transfer() else nullcontext() ) @@ -87,6 +74,5 @@ def _get_ec_connector_output( output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) ) - output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b8288c2d6481..43f5c453234c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -982,7 +982,6 @@ def __init__( # Ephemeral state transferred between execute_model() and sample_tokens(). self.execute_model_state: ExecuteModelState | None = None self.kv_connector_output: KVConnectorOutput | None = None - self.ec_connector_output: ECConnectorOutput | None = None self.mamba_state_idx: dict[str, int] = {} self._mamba_bufs: mamba_utils.MambaBuffers | None = None self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None @@ -3504,7 +3503,6 @@ def _pool( num_scheduled_tokens: int, num_scheduled_tokens_np: np.ndarray, kv_connector_output: KVConnectorOutput | None, - ec_connector_output: ECConnectorOutput | None, ) -> ModelRunnerOutput | AsyncModelRunnerOutput: num_reqs = self.input_batch.num_reqs assert num_reqs == len(self.input_batch.pooling_params), ( @@ -3542,7 +3540,6 @@ def _pool( req_ids=self.input_batch.req_ids.copy(), req_id_to_index=self.input_batch.req_id_to_index.copy(), kv_connector_output=kv_connector_output, - ec_connector_output=ec_connector_output, ) if raw_pooler_output is None or not any(finished_mask): @@ -4306,13 +4303,7 @@ def execute_model( encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) - # with_ec_conn_output tests is_empty(), so it must run after - # the context manager's finally block has populated - # ec_connector_worker_meta. - return ModelRunnerOutput.with_ec_conn_output( - make_empty_encoder_model_runner_output(scheduler_output), - ec_connector_output, - ) + return make_empty_encoder_model_runner_output(scheduler_output) if not num_scheduled_tokens: if ( @@ -4327,26 +4318,10 @@ def execute_model( # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. self._dummy_run(1) - if not has_kv_transfer_group() and not has_ec_transfer(): + if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT - output = ( - self.kv_connector_no_forward(scheduler_output, self.vllm_config) - if has_kv_transfer_group() - else EMPTY_MODEL_RUNNER_OUTPUT - ) - # EC transfer only ever runs on the first PP rank (that's - # where the multimodal encoder lives, see the is_first_rank - # gate above in _preprocess); other ranks have nothing of - # their own to report and must not touch encoder_cache. - if has_ec_transfer() and get_pp_group().is_first_rank: - ec_output = self.ec_connector_no_forward( - scheduler_output, self.encoder_cache - ) - output = ModelRunnerOutput.with_ec_conn_output( - output, ec_output.ec_connector_output - ) - return output + return self.kv_connector_no_forward(scheduler_output, self.vllm_config) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( @@ -4580,7 +4555,6 @@ def execute_model( # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) self.kv_connector_output = kv_connector_output - self.ec_connector_output = ec_connector_output return hidden_states if self.is_pooling_model: @@ -4590,7 +4564,6 @@ def execute_model( num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output, - ec_connector_output, ) sample_hidden_states = hidden_states[logits_indices] @@ -4668,16 +4641,12 @@ def sample_tokens( if self.execute_model_state is None: kv_connector_output = self.kv_connector_output self.kv_connector_output = None - ec_connector_output = self.ec_connector_output - self.ec_connector_output = None # receive sampled token ids from the last PP rank. if self.use_async_scheduling and not get_pp_group().is_last_rank: self._pp_receive_prev_sampled_token_ids_to_input_batch() - # In case of PP with kv/ec transfer, we need to pass through their - # outputs -- this rank never has a "real" ModelRunnerOutput of its - # own (see the is_last_rank early return in execute_model above). - output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) - return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) # Unpack ephemeral state. ( From 983f7b389f691af66ed82a86f26b824eadf2aac2 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Wed, 12 Aug 2026 12:06:20 +0300 Subject: [PATCH 26/31] Reduced lines of change to a minimum - reduced tests, reverted MRv1 changes, made docstrings more compact and functions leaner. Signed-off-by: omerpaz95 --- .../unit/test_ec_output_aggregator.py | 75 ++----- .../unit/test_worker_ec_connector.py | 203 ++++-------------- tests/v1/test_outputs.py | 22 +- .../ec_transfer/ec_connector/utils.py | 14 +- vllm/v1/executor/abstract.py | 1 - vllm/v1/executor/multiproc_executor.py | 19 +- vllm/v1/outputs.py | 17 +- vllm/v1/worker/gpu/model_runner.py | 17 +- 8 files changed, 89 insertions(+), 279 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py index d50647ca8840..e859e939c352 100644 --- a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -13,25 +13,15 @@ class FakeWorkerMeta(ECConnectorWorkerMetadata): - """Per-worker save/load reports, concatenated in merge order. - - `aggregate` returns a new object, as the base class declares: an - aggregator that discarded the return value would lose the merge. + """Records merge order. `aggregate` returns a new object, as the base class + declares: an aggregator discarding the return value would lose the merge. """ - def __init__( - self, - completed_saves: list[str] | None = None, - completed_loads: list[str] | None = None, - ): - self.completed_saves = completed_saves or [] - self.completed_loads = completed_loads or [] + def __init__(self, saves: list[str]): + self.saves = saves def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": - return FakeWorkerMeta( - self.completed_saves + other.completed_saves, - self.completed_loads + other.completed_loads, - ) + return FakeWorkerMeta(self.saves + other.saves) def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: @@ -40,55 +30,31 @@ def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: ) -def test_aggregate_unions_finished_ids_onto_output_rank(): - """EC work done on any rank reaches the scheduler via output_rank's output.""" - outputs = [ - _worker_output( - ECConnectorOutput(finished_sending={"mm0"}, finished_recving={"mm1"}) - ), - _worker_output(ECConnectorOutput(finished_sending={"mm2"})), - ] - - result = ECOutputAggregator().aggregate(outputs, output_rank=1) +def test_aggregate_folds_every_rank_onto_output_rank(): + """EC work done on any rank reaches the scheduler via output_rank's output. - assert result is outputs[1] - assert result.ec_connector_output.finished_sending == {"mm0", "mm2"} - assert result.ec_connector_output.finished_recving == {"mm1"} - - -def test_aggregate_worker_meta_folds_across_ranks(): - """Worker metadata is folded left across ranks, keeping each merge result.""" + The middle rank reports no worker metadata: it must neither seed nor clobber + the accumulator. + """ outputs = [ _worker_output( - ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"], [])) + ECConnectorOutput( + finished_sending={"mm0"}, + ec_connector_worker_meta=FakeWorkerMeta(["mm0"]), + ) ), + _worker_output(ECConnectorOutput(finished_recving={"mm1"})), _worker_output( - ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta([], ["mm1"])) - ), - _worker_output( - ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"], ["mm3"])) + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"])) ), ] result = ECOutputAggregator().aggregate(outputs, output_rank=2) - worker_meta = result.ec_connector_output.ec_connector_worker_meta - assert worker_meta.completed_saves == ["mm0", "mm2"] - assert worker_meta.completed_loads == ["mm1", "mm3"] - - -def test_aggregate_worker_meta_tolerates_ranks_without_meta(): - """Ranks reporting no metadata neither seed nor clobber the accumulator.""" - worker_meta = FakeWorkerMeta(["mm1"], []) - outputs = [ - _worker_output(ECConnectorOutput(finished_sending={"mm0"})), - _worker_output(ECConnectorOutput(ec_connector_worker_meta=worker_meta)), - _worker_output(ECConnectorOutput(finished_recving={"mm2"})), - ] - - result = ECOutputAggregator().aggregate(outputs, output_rank=0) - - assert result.ec_connector_output.ec_connector_worker_meta is worker_meta + assert result is outputs[2] + assert result.ec_connector_output.finished_sending == {"mm0"} + assert result.ec_connector_output.finished_recving == {"mm1"} + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"] def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): @@ -99,7 +65,6 @@ def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): assert result is outputs[0] assert result.ec_connector_output is None - assert ECOutputAggregator().aggregate([None], output_rank=0) is None diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py index 3834bdf60c9a..3dcad1e50ae5 100644 --- a/tests/v1/ec_connector/unit/test_worker_ec_connector.py +++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py @@ -3,198 +3,75 @@ """Unit tests for the V2 GPU model runner's EC connector wrapper.""" from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorBase, ECConnectorMetadata, - ECConnectorWorkerMetadata, ) from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT -from vllm.v1.worker.gpu.ec_connector import ( - NO_OP_EC_CONNECTOR, - ActiveECConnector, - get_ec_connector, -) +from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector pytestmark = pytest.mark.cpu_test +WORKER_META = object() -class FakeMetadata(ECConnectorMetadata): - pass - - -class FakeWorkerMeta(ECConnectorWorkerMetadata): - """Mirrors ECCPUWorkerMetadata: the transfers that completed this step.""" - - def __init__( - self, - completed_saves: list[str] | None = None, - completed_loads: list[str] | None = None, - ): - self.completed_saves = completed_saves or [] - self.completed_loads = completed_loads or [] - - def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": - return self - - -class FakeECConnector(ECConnectorBase): - """Records the calls that ActiveECConnector is expected to drive.""" - - def __init__(self, is_producer: bool = True, is_consumer: bool = False): - # ECConnectorBase.__init__ requires a full VllmConfig; set only what - # the worker-side code reads. - self._connector_metadata = None - self._is_producer = is_producer - self._is_consumer = is_consumer - self.bound_metadata: list[ECConnectorMetadata] = [] - self.clear_calls = 0 - self.load_calls: list[dict] = [] - self.saved_hashes: list[str] = [] - self.worker_meta_calls = 0 - self.worker_meta: FakeWorkerMeta | None = None - self.finished: tuple[set[str] | None, set[str] | None] = (None, None) - - def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> None: - self.bound_metadata.append(connector_metadata) - super().bind_connector_metadata(connector_metadata) - - def clear_connector_metadata(self) -> None: - self.clear_calls += 1 - super().clear_connector_metadata() - - def start_load_caches(self, encoder_cache: dict, **kwargs) -> None: - self.load_calls.append(encoder_cache) - - def save_caches(self, encoder_cache: dict, mm_hash: str) -> None: - self.saved_hashes.append(mm_hash) - - def get_finished(self, finished_req_ids: set[str]): - return self.finished - - def build_connector_worker_meta(self) -> FakeWorkerMeta | None: - self.worker_meta_calls += 1 - return self.worker_meta - - # Scheduler-side abstract methods, unused by these tests. - def has_cache_item(self, identifier: str) -> bool: - return False - - def update_state_after_alloc(self, request, index: int) -> None: - pass - def build_connector_meta(self, scheduler_output) -> ECConnectorMetadata: - return FakeMetadata() - - -def _scheduler_output(metadata: ECConnectorMetadata | None) -> SimpleNamespace: - return SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=frozenset()) +def _scheduler_output() -> SimpleNamespace: + return SimpleNamespace( + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset() + ) -def _active_connector( - encoder_cache: dict | None = None, **connector_kwargs -) -> tuple[ActiveECConnector, FakeECConnector]: - fake = FakeECConnector(**connector_kwargs) +def _connector( + encoder_cache: dict | None = None, + is_producer: bool = True, + is_consumer: bool = False, +) -> tuple[ActiveECConnector, MagicMock]: + fake = MagicMock(spec=ECConnectorBase) + fake.is_producer = is_producer + fake.is_consumer = is_consumer + fake.get_finished.return_value = (None, None) + fake.build_connector_worker_meta.return_value = WORKER_META with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): - connector = ActiveECConnector(SimpleNamespace(), encoder_cache or {}) - return connector, fake - - -def test_no_forward_is_noop_without_ec_connector(): - """`ec_connector_metadata is None` means no EC connector is configured: - nothing to poll, and the shared empty output is returned as-is. - """ - scheduler_output = _scheduler_output(metadata=None) - - assert NO_OP_EC_CONNECTOR.no_forward(scheduler_output) is EMPTY_MODEL_RUNNER_OUTPUT - - connector, fake = _active_connector() - assert connector.no_forward(scheduler_output) is EMPTY_MODEL_RUNNER_OUTPUT - assert fake.bound_metadata == [] - assert fake.worker_meta_calls == 0 - - -def test_no_forward_polls_connector_with_empty_metadata(): - """A step with no work still reaps completed transfers. - - The scheduler sends metadata every step, empty or not, and that is what - drives build_connector_worker_meta() -- the only channel by which finished - saves and loads are reported. - """ - connector, fake = _active_connector() - fake.worker_meta = FakeWorkerMeta(completed_saves=["mm0"]) - scheduler_output = _scheduler_output(metadata=FakeMetadata()) - - output = connector.no_forward(scheduler_output) - - assert fake.worker_meta_calls == 1 - assert output.ec_connector_output.ec_connector_worker_meta is fake.worker_meta - assert fake.bound_metadata == [scheduler_output.ec_connector_metadata] - assert fake.clear_calls == 1 + return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake @pytest.mark.parametrize( - ("is_producer", "is_consumer"), - [(True, False), (True, True), (False, True)], + ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)] ) -def test_maybe_get_output_saves_only_newly_added_caches(is_producer, is_consumer): - """Every producer offloads caches computed during the step, including an - ec_both node, and never re-offloads ones that were already cached. - """ +def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer): + """An ec_both node is also a producer: it must offload what it just computed.""" encoder_cache = {"mm_old": None} - connector, fake = _active_connector( - encoder_cache, is_producer=is_producer, is_consumer=is_consumer - ) + connector, fake = _connector(encoder_cache, is_producer, is_consumer) - with connector.maybe_get_output(_scheduler_output(FakeMetadata())): + with connector.maybe_get_output(_scheduler_output()): encoder_cache["mm_new"] = None - assert fake.saved_hashes == (["mm_new"] if is_producer else []) - assert fake.load_calls == ([encoder_cache] if is_consumer else []) + saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list] + assert saved == (["mm_new"] if is_producer else []) + assert fake.start_load_caches.called == is_consumer -def test_worker_meta_is_populated_only_on_exit(): - """The worker's report lands in the finally block, so a caller returning - from inside the `with` would drop it. - """ - connector, fake = _active_connector() - fake.worker_meta = FakeWorkerMeta(completed_loads=["mm0"]) +def test_worker_meta_is_reported_on_context_exit(): + """Reported in the finally block, so is_empty() sees it only after the exit.""" + connector, fake = _connector() - with connector.maybe_get_output(_scheduler_output(FakeMetadata())) as output: + with connector.maybe_get_output(_scheduler_output()) as output: assert output.ec_connector_worker_meta is None - assert fake.worker_meta_calls == 0 - assert output.ec_connector_worker_meta is fake.worker_meta + assert output.ec_connector_worker_meta is WORKER_META + assert fake.clear_connector_metadata.called -def test_get_ec_connector_activates_only_for_a_multimodal_ec_deployment(): - encoder_cache = SimpleNamespace(encoder_outputs={}) - config = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False)) - encoder_decoder = SimpleNamespace( - model_config=SimpleNamespace(is_encoder_decoder=True) - ) +def test_no_forward_reports_without_running_the_model(): + connector, _ = _connector() + + output = connector.no_forward(_scheduler_output()) + + assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META - with patch( - "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", - return_value=FakeECConnector(), - ): - with patch( - "vllm.v1.worker.gpu.ec_connector.has_ec_transfer", return_value=False - ): - assert get_ec_connector(config, encoder_cache) is NO_OP_EC_CONNECTOR - - with patch( - "vllm.v1.worker.gpu.ec_connector.has_ec_transfer", return_value=True - ): - assert get_ec_connector(encoder_decoder, encoder_cache) is ( - NO_OP_EC_CONNECTOR - ) - assert get_ec_connector(config, None) is NO_OP_EC_CONNECTOR - connector = get_ec_connector(config, encoder_cache) - - assert isinstance(connector, ActiveECConnector) - # Aliases the live dict, so caches added later in the step are visible. - assert connector.encoder_cache is encoder_cache.encoder_outputs + empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output()) + assert empty is EMPTY_MODEL_RUNNER_OUTPUT diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index c9322fd29409..b7f1ae579b9d 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -130,31 +130,19 @@ def test_slice_all_requests(self): assert sliced.cu_num_generated_tokens is None -def _ec_output() -> ECConnectorOutput: - return ECConnectorOutput(finished_sending={"mm_hash"}) - - -def test_with_ec_conn_output_keeps_output_when_nothing_to_report(): +def test_with_ec_conn_output_sets_field_in_place(): output = ModelRunnerOutput(req_ids=["r0"], req_id_to_index={"r0": 0}) + ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) - assert ModelRunnerOutput.with_ec_conn_output(output, None) is output assert ModelRunnerOutput.with_ec_conn_output(output, ECConnectorOutput()) is output assert output.ec_connector_output is None - - -def test_with_ec_conn_output_sets_field_in_place(): - output = ModelRunnerOutput(req_ids=["r0"], req_id_to_index={"r0": 0}) - ec_output = _ec_output() - - result = ModelRunnerOutput.with_ec_conn_output(output, ec_output) - - assert result is output - assert result.ec_connector_output is ec_output + assert ModelRunnerOutput.with_ec_conn_output(output, ec_output) is output + assert output.ec_connector_output is ec_output def test_with_ec_conn_output_copies_shared_empty_output(): """The shared empty output is copied, never written to.""" - ec_output = _ec_output() + ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output) diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index e9cd57115dda..41fca433f945 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -6,11 +6,11 @@ class ECOutputAggregator: - """Merge every worker's EC connector output into the one ModelRunnerOutput - that reaches the scheduler. + """Merge every worker's EC connector output onto the single + ModelRunnerOutput that reaches the scheduler. Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the - scheduler, but the EC connector may have run on other ranks. + scheduler, but the EC connector may have run on any rank. """ def aggregate( @@ -32,10 +32,10 @@ def aggregate( finished_sending |= ec_output.finished_sending or set() finished_recving |= ec_output.finished_recving or set() - if worker_meta is None: - worker_meta = ec_output.ec_connector_worker_meta - elif other := ec_output.ec_connector_worker_meta: - worker_meta = worker_meta.aggregate(other) + if meta := ec_output.ec_connector_worker_meta: + worker_meta = ( + meta if worker_meta is None else worker_meta.aggregate(meta) + ) aggregated = ECConnectorOutput( finished_sending=finished_sending or None, diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 0b570e2f38a6..9a0f86c0b74a 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -286,7 +286,6 @@ def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: ) def init_ec_output_aggregator(self) -> None: - """Init ECOutputAggregator""" self.ec_output_aggregator = ECOutputAggregator() @cached_property # Avoid unnecessary RPC calls diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 0765f9b7501d..df20a90d51d3 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -381,8 +381,7 @@ def collective_rpc( # type: ignore[override] ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: """Returns single result if unique_reply_rank and/or an output - aggregator (kv_output_aggregator/ec_output_aggregator) is provided, - otherwise list.""" + aggregator is provided, otherwise list.""" assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -392,22 +391,18 @@ def collective_rpc( # type: ignore[override] deadline = None if timeout is None else time.monotonic() + timeout kwargs = kwargs or {} - aggregators = [ - agg - for agg in (kv_output_aggregator, ec_output_aggregator) - if agg is not None - ] + aggregators = [a for a in (kv_output_aggregator, ec_output_aggregator) if a] aggregate: Callable[[Any], Any] if aggregators: output_rank = None def _aggregate(outputs: Any) -> Any: # Each aggregator merges its own connector's output onto - # outputs[output_rank] in place and returns it, so chaining - # them and keeping the last result is safe. - result = None - for agg in aggregators: - result = agg.aggregate(outputs, output_rank=unique_reply_rank or 0) + # outputs[output_rank] in place and returns it, so chaining is safe. + rank = unique_reply_rank or 0 + result = outputs[rank] + for a in aggregators: + result = a.aggregate(outputs, output_rank=rank) return result aggregate = _aggregate diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 564a80e41731..9bd47eb444b5 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -334,14 +334,10 @@ def with_kv_conn_output_only( def with_ec_conn_output_only( ec_connector_output: ECConnectorOutput | None, ) -> "ModelRunnerOutput": - """Return ModelRunnerOutput containing the provided ECConnectorOutput, - otherwise empty. - """ - if ec_connector_output is None or ec_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.ec_connector_output = ec_connector_output - return output + """Return an otherwise-empty output carrying `ec_connector_output`.""" + return ModelRunnerOutput.with_ec_conn_output( + EMPTY_MODEL_RUNNER_OUTPUT, ec_connector_output + ) @staticmethod def with_ec_conn_output( @@ -350,9 +346,8 @@ def with_ec_conn_output( ) -> "ModelRunnerOutput": """Return `output` carrying `ec_connector_output`. - Sets the field on `output` and returns it, except that the shared - empty output is copied first rather than written to. Callers must use - the return value. + The shared empty output is copied rather than written to, so callers + must use the return value. """ if ec_connector_output is None or ec_connector_output.is_empty(): return output diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0a98a61f4883..ec86855914e1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1248,12 +1248,7 @@ def postprocess_sampled( def _merge_ec_connector_no_forward( self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput ) -> ModelRunnerOutput: - """Merge the EC connector's output into `output` for a - step with no work to run. - - A no-op unless this rank runs the EC connector: the connector is the - no-op one unless an encoder cache exists. - """ + """Let the EC connector send/recv on a step with no work to run.""" return ModelRunnerOutput.with_ec_conn_output( output, self.ec_connector.no_forward(scheduler_output).ec_connector_output, @@ -1568,9 +1563,7 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - # This rank never has a "real" ModelRunnerOutput of its own (see the - # is_last_pp_rank early return in execute_model above), but may have - # produced ec_connector_output on the first PP rank -- pass it through. + # The first PP rank holds the encoder cache, so pass its EC output on. output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) @@ -1709,10 +1702,8 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: if not self.is_last_pp_rank: self.postprocess_num_computed_tokens(input_batch) - return ModelRunnerOutput.with_ec_conn_output( - ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output), - ec_connector_output, - ) + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) assert self.pooling_runner is not None pooler_output, finished_mask = self.pooling_runner.pool( From 753a10f33d04fdb51ba99494b7a2e4495547de94 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Thu, 13 Aug 2026 10:46:31 +0300 Subject: [PATCH 27/31] Added test + fix for not mutating global EMPTY_MODEL_RUNNER_OUTPUT singleton. Also explicitly reject legacy Ray backend. Signed-off-by: omerpaz95 --- .../unit/test_ec_output_aggregator.py | 27 ++++++++++++++++++- .../ec_transfer/ec_connector/utils.py | 9 +++++-- vllm/v1/executor/ray_executor.py | 13 +++++++++ 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py index e859e939c352..4d5b7dae1037 100644 --- a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -7,7 +7,12 @@ from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator -from vllm.v1.outputs import ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + KVConnectorOutput, + ModelRunnerOutput, +) pytestmark = pytest.mark.cpu_test @@ -68,6 +73,26 @@ def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): assert ECOutputAggregator().aggregate([None], output_rank=0) is None +def test_aggregate_does_not_write_through_the_shared_empty_output(): + """A rank with nothing to report yields the shared empty output singleton. + + Folding another rank's metadata onto it must not write through to the + module-level object, which every later step would then carry. + """ + outputs = [ + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"])) + ), + EMPTY_MODEL_RUNNER_OUTPUT, + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=1) + + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"] + + def test_chaining_with_kv_aggregator_preserves_both_outputs(): """MultiprocExecutor chains both aggregators and keeps only the last result, so each must merge onto the same output_rank output rather than replace it. diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index 41fca433f945..f5f78e6c3c33 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -42,5 +42,10 @@ def aggregate( finished_recving=finished_recving or None, ec_connector_worker_meta=worker_meta, ) - output.ec_connector_output = None if aggregated.is_empty() else aggregated - return output + if aggregated.is_empty(): + output.ec_connector_output = None + return output + + # `output` is the shared empty output whenever `output_rank` had no work, + # so attach through the copy-on-write helper. + return ModelRunnerOutput.with_ec_conn_output(output, aggregated) diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 39749ffc257e..986e4d9bbb03 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -89,6 +89,19 @@ def _init_executor(self) -> None: # KV connector setup self.has_connector = self.vllm_config.kv_transfer_config is not None + if ( + self.vllm_config.ec_transfer_config is not None + and self.parallel_config.world_size > 1 + ): + raise NotImplementedError( + "EC connector worker metadata is not supported with the " + "legacy Ray executor when world_size > 1: only the output " + "of a single worker is fetched, silently dropping the " + "other workers' EC connector state. Set " + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, " + "or use the multiprocessing executor instead." + ) + self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and ( self.vllm_config.ec_transfer_config is None or self.vllm_config.ec_transfer_config.is_ec_consumer From 7d898df1d1c6e8a47571e8221de97a9b704aabfe Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Thu, 13 Aug 2026 19:03:26 +0300 Subject: [PATCH 28/31] Update tests/v1/test_outputs.py Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Signed-off-by: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> --- tests/v1/test_outputs.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 1b9da2185eaa..883a576e5197 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -203,14 +203,6 @@ def test_slice_all_requests(self): assert sliced.cu_num_generated_tokens is None -def test_with_ec_conn_output_sets_field_in_place(): - output = ModelRunnerOutput(req_ids=["r0"], req_id_to_index={"r0": 0}) - ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) - - assert ModelRunnerOutput.with_ec_conn_output(output, ECConnectorOutput()) is output - assert output.ec_connector_output is None - assert ModelRunnerOutput.with_ec_conn_output(output, ec_output) is output - assert output.ec_connector_output is ec_output def test_with_ec_conn_output_copies_shared_empty_output(): From 9ac0d01b78944e9481495f40d47e9ec0b1d9ccf6 Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Thu, 13 Aug 2026 19:28:30 +0300 Subject: [PATCH 29/31] removed a test. Signed-off-by: omerpaz95 --- tests/v1/test_outputs.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 883a576e5197..445b5a936da6 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -203,8 +203,6 @@ def test_slice_all_requests(self): assert sliced.cu_num_generated_tokens is None - - def test_with_ec_conn_output_copies_shared_empty_output(): """The shared empty output is copied, never written to.""" ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) From b0004d24ab674c232b70af630b824187ee53227a Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Fri, 14 Aug 2026 13:42:43 +0300 Subject: [PATCH 30/31] Fix tests to pass ci. Signed-off-by: omerpaz95 --- tests/v1/executor/test_executor.py | 3 +++ tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py | 1 + 2 files changed, 4 insertions(+) diff --git a/tests/v1/executor/test_executor.py b/tests/v1/executor/test_executor.py index c529c3204d50..a21a8dafd328 100644 --- a/tests/v1/executor/test_executor.py +++ b/tests/v1/executor/test_executor.py @@ -9,6 +9,7 @@ import pytest +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.sampling_params import SamplingParams @@ -98,6 +99,7 @@ def collective_rpc( non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any | list[Any] | Future[Any | list[Any]]: # Drop marker to show that this was run with open(".marker", "w"): @@ -110,6 +112,7 @@ def collective_rpc( non_block, unique_reply_rank, kv_output_aggregator, + ec_output_aggregator, ) diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index 0c0f9f1f8998..e5cbd977b6d8 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -55,6 +55,7 @@ def _run_engine_core_handshake( class _FakeScheduler: def __init__(self, **kwargs: Any) -> None: self.connector = connector + self.ec_connector = None def get_kv_connector(self) -> KVConnectorBase_V1: return connector From 2cf705c8a108dd09833ffd50b7a6d40bacb597cc Mon Sep 17 00:00:00 2001 From: omerpaz95 Date: Fri, 14 Aug 2026 17:45:57 +0300 Subject: [PATCH 31/31] Fix eplb CI Signed-off-by: omerpaz95 --- tests/v1/worker/test_gpu_model_runner_v2_eplb.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index fbeaa198d0d6..ebb4beb2a5a9 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -180,6 +180,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): hidden_states=None, aux_hidden_states=None, finished_req_ids=set(), + ec_connector_output=None, routed_experts=None, num_tokens_across_dp=None, )