[Bugfix][EPD] Skip sampling for encoder-only async steps (#58490)
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Tianyu Guo committed
Sep 28, 2026 at 11:35 UTC
53d2e16a9779b275a5bb98b69fba61543581a2ed
3 files changed
+36
-5
tests/v1/engine/test_engine_core.py
+33
@@ -4,6 +4,7 @@
4
import copy
5
import time
6
import uuid
7
+from collections import deque
8
from concurrent.futures import Future, ThreadPoolExecutor
9
from unittest.mock import MagicMock, PropertyMock, patch
10
@@ -395,6 +396,38 @@ def test_engine_core_concurrent_batches():
396
req_id = (req_id + 1) % 2
397
398
399
+@pytest.mark.parametrize("encoder_only", [True, False])
400
+def test_async_step_skips_sampling_for_mm_encoder_only(encoder_only):
401
+ """Encoder outputs are complete; only language-model steps need sampling."""
402
+ core = MagicMock()
403
+ core.batch_queue = deque()
404
+ core.batch_queue_size = 2
405
+ core.is_mm_encoder_only = encoder_only
406
+ core.is_pooling_model = False
407
+ core.scheduler.has_requests.side_effect = [True, False, False]
408
+ scheduled = core.scheduler.schedule.return_value
409
+ scheduled.total_num_scheduled_tokens = 16
410
+ scheduled.pending_structured_output_tokens = False
411
+
412
+ output = MagicMock()
413
+ executed: Future[ModelRunnerOutput | None] = Future()
414
+ executed.set_result(output if encoder_only else None)
415
+ sampled: Future[ModelRunnerOutput | None] = Future()
416
+ sampled.set_result(None if encoder_only else output)
417
+ core.model_executor.execute_model.return_value = executed
418
+ core.model_executor.sample_tokens.return_value = sampled
419
+
420
+ EngineCore.step_with_batch_queue(core)
421
+ if core.batch_queue:
422
+ EngineCore.step_with_batch_queue(core)
423
+
424
+ core.scheduler.update_from_output.assert_called_once_with(scheduled, output)
425
+ if encoder_only:
426
+ core.model_executor.sample_tokens.assert_not_called()
427
+ else:
428
+ core.model_executor.sample_tokens.assert_called_once()
429
+
430
+
431
@multi_gpu_test(num_gpus=2)
432
def test_engine_core_tp():
433
"""Test engine can initialize worker in tp properly."""
tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
+1
@@ -99,6 +99,7 @@ def _run_engine_core_handshake(
99
),
100
speculative_config=None,
101
ec_transfer_config=None,
102
+ is_mm_encoder_only=False,
103
max_concurrent_batches=1,
104
model_config=SimpleNamespace(
105
runner_type="generate",
vllm/v1/engine/core.py
+2
-5
@@ -219,10 +219,7 @@ class EngineCore:
219
logger.debug("Batch queue is enabled with size %d", self.batch_queue_size)
220
self.batch_queue = deque(maxlen=self.batch_queue_size)
221
222
- self.is_ec_consumer = (
223
- vllm_config.ec_transfer_config is None
224
- or vllm_config.ec_transfer_config.is_ec_consumer
225
- )
222
+ self.is_mm_encoder_only = vllm_config.is_mm_encoder_only
223
self.is_pooling_model = vllm_config.model_config.runner_type == "pooling"
224
225
self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None
@@ -702,7 +699,7 @@ class EngineCore:
699
exec_future = self.model_executor.execute_model(
700
scheduler_output, non_block=True
701
)
705
- if self.is_ec_consumer:
702
+ if not self.is_mm_encoder_only:
703
model_executed = scheduler_output.total_num_scheduled_tokens > 0
704
705
if self.is_pooling_model or not model_executed: