From cab5c9a2a9601ce27bd765db7158f6fce6c73fdf Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 2 Jun 2026 08:57:25 -0700 Subject: [PATCH] [Core] Move `max_concurrent_batches` to `VllmConfig` (#44274) Signed-off-by: Nick Hill --- tests/distributed/test_multiproc_executor.py | 2 +- tests/distributed/test_ray_v2_executor.py | 2 +- .../model_loader/tensorizer_loader/conftest.py | 4 ---- tests/v1/engine/test_engine_core.py | 17 ++++++++++++----- tests/v1/engine/test_engine_core_client.py | 1 - vllm/config/vllm.py | 9 +++++++++ vllm/v1/engine/core.py | 2 +- vllm/v1/executor/abstract.py | 4 ---- vllm/v1/executor/multiproc_executor.py | 8 +------- vllm/v1/executor/ray_executor.py | 8 -------- vllm/v1/executor/uniproc_executor.py | 5 ----- 11 files changed, 25 insertions(+), 37 deletions(-) diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 29d7f94c510..20dd4f36393 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -284,7 +284,7 @@ def test_multiproc_executor_pipeline_parallel(): assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)" # Verify max_concurrent_batches for pipeline parallel - assert executor.max_concurrent_batches == 2, ( + assert vllm_config.max_concurrent_batches == 2, ( "Max concurrent batches should equal PP size" ) diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py index 5daec22df6f..398ee30c068 100644 --- a/tests/distributed/test_ray_v2_executor.py +++ b/tests/distributed/test_ray_v2_executor.py @@ -83,7 +83,7 @@ def assert_executor(executor, tp_size, pp_size): assert executor._get_output_rank() == expected_output_rank if pp_size > 1: - assert executor.max_concurrent_batches == pp_size + assert executor.vllm_config.max_concurrent_batches == pp_size executor.check_health() assert not executor.is_failed diff --git a/tests/model_executor/model_loader/tensorizer_loader/conftest.py b/tests/model_executor/model_loader/tensorizer_loader/conftest.py index 6c85a139919..051890e89a1 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/conftest.py +++ b/tests/model_executor/model_loader/tensorizer_loader/conftest.py @@ -87,10 +87,6 @@ class DummyExecutor(UniProcExecutor): self.collective_rpc("init_worker", args=([kwargs],)) self.collective_rpc("init_device") - @property - def max_concurrent_batches(self) -> int: - return 2 - def shutdown(self): if hasattr(self, "thread_pool"): self.thread_pool.shutdown(wait=False) diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index ae674919ae9..aa2a70559dd 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -5,6 +5,7 @@ import copy import time import uuid from concurrent.futures import Future, ThreadPoolExecutor +from unittest.mock import PropertyMock, patch import pytest from transformers import AutoTokenizer @@ -293,10 +294,6 @@ def test_engine_core_concurrent_batches(): # Use the thread pool instead of creating a new thread return self.thread_pool.submit(_execute) - @property - def max_concurrent_batches(self) -> int: - return 2 - def shutdown(self): if hasattr(self, "thread_pool"): self.thread_pool.shutdown(wait=False) @@ -314,7 +311,17 @@ def test_engine_core_concurrent_batches(): async_scheduling=False, ) vllm_config = engine_args.create_engine_config() - with set_default_torch_num_threads(1): + # Force two concurrent batches to exercise the batch queue independently + # of async scheduling (which is disabled above). + with ( + set_default_torch_num_threads(1), + patch.object( + VllmConfig, + "max_concurrent_batches", + new_callable=PropertyMock, + return_value=2, + ), + ): engine_core = EngineCore( vllm_config=vllm_config, log_stats=False, executor_class=DummyExecutor ) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 8327b5b0763..36dc95eea49 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -1212,7 +1212,6 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}] mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024] mock_executor.initialize_from_config.return_value = None - mock_executor.max_concurrent_batches = 1 return mock_executor diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f753647081c..db139b6532a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -487,6 +487,15 @@ class VllmConfig: ] return hash_str + @property + def max_concurrent_batches(self) -> int: + # PP requires PP-size concurrent batches to fill the pipeline. + # Async scheduling requires 2 concurrent batches to overlap. + pp_size = self.parallel_config.pipeline_parallel_size + if pp_size > 1: + return pp_size + return 2 if self.scheduler_config.async_scheduling else 1 + @property def num_speculative_tokens(self) -> int: if ( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d08bbc951fa..60583ea5887 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -188,7 +188,7 @@ class EngineCore: # Batch queue for scheduled batches. This enables us to asynchronously # schedule and execute batches, and is required by pipeline parallelism # to eliminate pipeline bubbles. - self.batch_queue_size = self.model_executor.max_concurrent_batches + self.batch_queue_size = vllm_config.max_concurrent_batches self.batch_queue: ( deque[tuple[Future[ModelRunnerOutput], SchedulerOutput, Future[Any]]] | None ) = None diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index e68c0283f57..7beef598e27 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -253,10 +253,6 @@ class Executor(ABC): output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") return output[0] - @property - def max_concurrent_batches(self) -> int: - return 1 - def profile(self, is_start: bool = True, profile_prefix: str | None = None): self.collective_rpc("profile", args=(is_start, profile_prefix)) diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 449fbdd9736..c5766c923c8 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -15,7 +15,7 @@ from concurrent.futures import Future, InvalidStateError from contextlib import suppress from dataclasses import dataclass from enum import Enum, auto -from functools import cached_property, partial +from functools import partial from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess from multiprocessing.synchronize import Lock as LockType @@ -472,12 +472,6 @@ class MultiprocExecutor(Executor): self.collective_rpc("check_health", timeout=10) return - @cached_property - def max_concurrent_batches(self) -> int: - # PP requires PP-size concurrent batches to fill the pipeline. - pp_size = self.parallel_config.pipeline_parallel_size - return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size - def _get_output_rank(self) -> int: # Only returns ModelRunnerOutput from TP rank=0 and PP rank=-1 # (the first TP worker of the last PP stage). diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index cfeebb5e09d..749e59e04c2 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -96,14 +96,6 @@ class RayDistributedExecutor(Executor): self.scheduler_output: SchedulerOutput | None = None - @property - def max_concurrent_batches(self) -> int: - """Ray distributed executor supports pipeline parallelism, - meaning that it allows PP size batches to be executed concurrently. - """ - pp_size = self.parallel_config.pipeline_parallel_size - return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size - def shutdown(self) -> None: if logger: # Somehow logger can be None here. diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py index c3be3300fd3..dd04b718d67 100644 --- a/vllm/v1/executor/uniproc_executor.py +++ b/vllm/v1/executor/uniproc_executor.py @@ -3,7 +3,6 @@ import os from collections.abc import Callable from concurrent.futures import Future -from functools import cached_property from multiprocessing import Lock from typing import Any @@ -77,10 +76,6 @@ class UniProcExecutor(Executor): local_rank = int(device_info[1]) if len(device_info) > 1 else 0 return distributed_init_method, 0, local_rank - @cached_property - def max_concurrent_batches(self) -> int: - return 2 if self.scheduler_config.async_scheduling else 1 - def collective_rpc( # type: ignore[override] self, method: str | Callable,