A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[P/D]`kv_output_aggregator` support P TP > D TP (#23917)
Signed-off-by: LCAIZJ <leichao139636@163.com> Co-authored-by: leichao.lc <leichao.lc@antgroup.com>
Chao Lei committed
Sep 15, 2025 at 17:36 UTC
8de261b04a0a0e916d3d25d528d0f2ddeede2a6b
5 files changed
+21
-5
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+11
@@ -355,3 +355,14 @@ class KVConnectorBase_V1(ABC):
355
raise TypeError("get_required_kvcache_layout should not be called "
356
"on the abstract base class")
357
return None
358
+
359
+ def get_finished_count(self) -> Optional[int]:
360
+ """
361
+ Get the count of requests expected to complete send/receive operations
362
+ via this connector.
363
+
364
+ Returns:
365
+ int: expected sending or receiving completion count.
366
+ """
367
+
368
+ return None
\ No newline at end of file
vllm/executor/executor_base.py
+7
@@ -13,6 +13,7 @@ from typing_extensions import TypeVar
13
14
import vllm.platforms
15
from vllm.config import VllmConfig
16
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
17
from vllm.logger import init_logger
18
from vllm.lora.request import LoRARequest
19
from vllm.model_executor.layers.sampler import SamplerOutput
@@ -54,6 +55,7 @@ class ExecutorBase(ABC):
55
self._init_executor()
56
self.is_sleeping = False
57
self.sleeping_tags: set[str] = set()
58
+ self.kv_output_aggregator = None
59
60
@abstractmethod
61
def _init_executor(self) -> None:
@@ -252,6 +254,11 @@ class ExecutorBase(ABC):
254
exception."""
255
self.check_health()
256
257
+ def init_kv_output_aggregator(self, finished_count: Optional[int]) -> None:
258
+ """Init KVOutputAggregator"""
259
+ self.kv_output_aggregator = KVOutputAggregator(
260
+ finished_count or self.parallel_config.world_size)
261
+
262
263
class DistributedExecutorBase(ExecutorBase):
264
"""Abstract superclass of distributed executor implementations."""
vllm/v1/engine/core.py
+3
@@ -128,6 +128,9 @@ class EngineCore:
128
log_stats=self.log_stats,
129
)
130
self.use_spec_decode = vllm_config.speculative_config is not None
131
+ if self.scheduler.connector is not None: # type: ignore
132
+ self.model_executor.init_kv_output_aggregator(
133
+ self.scheduler.connector.get_finished_count()) # type: ignore
134
135
self.mm_registry = mm_registry = MULTIMODAL_REGISTRY
136
self.mm_receiver_cache = engine_receiver_cache_from_config(
vllm/v1/executor/multiproc_executor.py
-3
@@ -26,7 +26,6 @@ from vllm.distributed import (destroy_distributed_environment,
26
destroy_model_parallel)
27
from vllm.distributed.device_communicators.shm_broadcast import (Handle,
28
MessageQueue)
29
-from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
29
from vllm.distributed.parallel_state import (get_dp_group, get_ep_group,
30
get_pp_group, get_tp_group)
31
from vllm.executor.multiproc_worker_utils import (
@@ -135,8 +134,6 @@ class MultiprocExecutor(Executor):
134
135
self.output_rank = self._get_output_rank()
136
self.has_connector = self.vllm_config.kv_transfer_config is not None
138
- self.kv_output_aggregator = KVOutputAggregator(
139
- self.parallel_config.world_size)
137
138
def start_worker_monitor(self):
139
workers = self.workers
vllm/v1/executor/ray_distributed_executor.py
-2
@@ -51,8 +51,6 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
51
52
# KV connector setup
53
self.has_connector = self.vllm_config.kv_transfer_config is not None
54
- self.kv_output_aggregator = KVOutputAggregator(
55
- self.parallel_config.world_size)
54
55
@property
56
def max_concurrent_batches(self) -> int: