From 61c1dd0966fd8d91954003cf6b7ab746448a46f8 Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Mon, 10 Aug 2026 09:25:57 +0300 Subject: [PATCH] [EC Connector] Call to EC Connector update_connector_output from scheduler (#49579) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: omerpaz95 Signed-off-by: Nicolò Lucchesi Co-authored-by: Nicolò Lucchesi Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/core/test_async_scheduler.py | 1 + tests/v1/core/test_scheduler.py | 37 +++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 5 ++++ 3 files changed, 43 insertions(+) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 1b01d184dc..79c6f6f65c 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -272,6 +272,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.perf_metrics = None scheduler.connector = None + scheduler.ec_connector = None scheduler.structured_output_manager = Mock() scheduler.structured_output_manager.should_advance.return_value = True scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 75f9402394..c419d50b5b 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -40,6 +40,7 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.outputs import ( DraftTokenIds, + ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, make_empty_encoder_model_runner_output, @@ -3237,6 +3238,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.perf_metrics = None scheduler.connector = None + scheduler.ec_connector = None scheduler.structured_output_manager = Mock() scheduler.structured_output_manager.should_advance.return_value = True scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( @@ -4771,6 +4773,41 @@ def test_scheduler_kv_connector_stats(): assert final_stats == expected_data +def test_ec_connector_update_connector_output_called(): + """Test that worker-side EC connector output is forwarded to the + EC connector's update_connector_output hook.""" + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + use_ec_connector=True, + ec_role="ec_consumer", + ) + scheduler.ec_connector.update_connector_output = Mock() + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + ec_connector_output = ECConnectorOutput(finished_sending={"hash_test1"}) + model_runner_output = ModelRunnerOutput( + req_ids=[], + req_id_to_index={}, + ec_connector_output=ec_connector_output, + ) + + scheduler.update_from_output(scheduler_output, model_runner_output) + + scheduler.ec_connector.update_connector_output.assert_called_once_with( + ec_connector_output + ) + + # ============================================================================== # Variable-length encoder cross-attention block allocation tests # ============================================================================== diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 006da8c7d1..15ef101799 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1672,6 +1672,7 @@ class Scheduler(SchedulerInterface): pooler_outputs = model_runner_output.pooler_output num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + ec_connector_output = model_runner_output.ec_connector_output cudagraph_stats = model_runner_output.cudagraph_stats # Every GPU write enqueued by this and earlier steps has completed, so it is @@ -1979,6 +1980,10 @@ class Scheduler(SchedulerInterface): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # EC Connector: update state from worker-side EC connector output. + if self.ec_connector is not None and ec_connector_output: + self.ec_connector.update_connector_output(ec_connector_output) + # Worker-side KV connector stats from the model runner output. kv_connector_stats: KVConnectorStats | None = ( kv_connector_output.kv_connector_stats if kv_connector_output else None