diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index e16e0550c2..2d9f7e8632 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -4,6 +4,7 @@ import copy import time import uuid +from collections import deque from concurrent.futures import Future, ThreadPoolExecutor from unittest.mock import MagicMock, PropertyMock, patch @@ -395,6 +396,38 @@ def test_engine_core_concurrent_batches(): req_id = (req_id + 1) % 2 +@pytest.mark.parametrize("encoder_only", [True, False]) +def test_async_step_skips_sampling_for_mm_encoder_only(encoder_only): + """Encoder outputs are complete; only language-model steps need sampling.""" + core = MagicMock() + core.batch_queue = deque() + core.batch_queue_size = 2 + core.is_mm_encoder_only = encoder_only + core.is_pooling_model = False + core.scheduler.has_requests.side_effect = [True, False, False] + scheduled = core.scheduler.schedule.return_value + scheduled.total_num_scheduled_tokens = 16 + scheduled.pending_structured_output_tokens = False + + output = MagicMock() + executed: Future[ModelRunnerOutput | None] = Future() + executed.set_result(output if encoder_only else None) + sampled: Future[ModelRunnerOutput | None] = Future() + sampled.set_result(None if encoder_only else output) + core.model_executor.execute_model.return_value = executed + core.model_executor.sample_tokens.return_value = sampled + + EngineCore.step_with_batch_queue(core) + if core.batch_queue: + EngineCore.step_with_batch_queue(core) + + core.scheduler.update_from_output.assert_called_once_with(scheduled, output) + if encoder_only: + core.model_executor.sample_tokens.assert_not_called() + else: + core.model_executor.sample_tokens.assert_called_once() + + @multi_gpu_test(num_gpus=2) def test_engine_core_tp(): """Test engine can initialize worker in tp properly.""" 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 509d3d3234..045f6cb425 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -99,6 +99,7 @@ def _run_engine_core_handshake( ), speculative_config=None, ec_transfer_config=None, + is_mm_encoder_only=False, max_concurrent_batches=1, model_config=SimpleNamespace( runner_type="generate", diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b36caf092c..5c09a0059e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -219,10 +219,7 @@ class EngineCore: logger.debug("Batch queue is enabled with size %d", self.batch_queue_size) self.batch_queue = deque(maxlen=self.batch_queue_size) - self.is_ec_consumer = ( - vllm_config.ec_transfer_config is None - or vllm_config.ec_transfer_config.is_ec_consumer - ) + self.is_mm_encoder_only = vllm_config.is_mm_encoder_only self.is_pooling_model = vllm_config.model_config.runner_type == "pooling" self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None @@ -702,7 +699,7 @@ class EngineCore: exec_future = self.model_executor.execute_model( scheduler_output, non_block=True ) - if self.is_ec_consumer: + if not self.is_mm_encoder_only: model_executed = scheduler_output.total_num_scheduled_tokens > 0 if self.is_pooling_model or not model_executed: