[Core] Move max_concurrent_batches to VllmConfig (#44274)

Signed-off-by: Nick Hill <[email protected]>

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"
         )
 

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

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)

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
         )

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
 

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 (

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

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))
 

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).

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.

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,