@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4dfdb821c8

Fork 0

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] Dynamic `kv_output_aggregator` collect size (#26734)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Oct 22, 2025 at 18:07 UTC 4dfdb821c8e46781ba6a37d828a7c76bcd401ae7

7 files changed +90 -19

tests/v1/kv_connector/unit/test_nixl_connector.py

+2 -2

@@ -703,7 +703,7 @@ def test_kv_connector_stats_aggregation():

703

704 # Create KVOutputAggregator for 3 workers (simulating TP=3), same thing

705 # done in MultiprocExecutor.execute_model

706 - aggregator = KVOutputAggregator(world_size=3)

706 + aggregator = KVOutputAggregator(expected_finished_count=3)

707

708 # Create stats for multiple workers with different transfer patterns

709 worker1_stats = NixlKVConnectorStats()

@@ -768,7 +768,7 @@ def test_multi_kv_connector_stats_aggregation():

768 KVOutputAggregator (used by MultiprocExecutor).

769 """

770

771 - aggregator = KVOutputAggregator(world_size=3)

771 + aggregator = KVOutputAggregator(expected_finished_count=3)

772

773 from dataclasses import dataclass

774

tests/v1/kv_connector/unit/test_output_aggregator.py renamed

+41 -2

@@ -16,11 +16,13 @@ class DummyModelRunnerOutput(ModelRunnerOutput):

16 finished_sending: set[str] | None = None,

17 finished_recving: set[str] | None = None,

18 invalid_block_ids: set[int] | None = None,

19 + expected_finished_count: int = 0,

20 ):

21 self.kv_connector_output = KVConnectorOutput(

22 finished_sending=finished_sending,

23 finished_recving=finished_recving,

24 invalid_block_ids=invalid_block_ids or set(),

25 + expected_finished_count=expected_finished_count,

26 )

28 def __repr__(self):

@@ -33,7 +35,7 @@ class DummyModelRunnerOutput(ModelRunnerOutput):

37 def test_aggregate_workers_output():

36 - aggregator = KVOutputAggregator(world_size=2)

38 + aggregator = KVOutputAggregator(expected_finished_count=2)

40 output1 = DummyModelRunnerOutput()

41 output2 = DummyModelRunnerOutput()

@@ -85,7 +87,7 @@ def test_aggregate_workers_output():

89 def test_async_aggregate_workers_output():

88 - aggregator = KVOutputAggregator(world_size=2)

90 + aggregator = KVOutputAggregator(expected_finished_count=2)

92 future1: Future[DummyModelRunnerOutput] = Future()

93 future2: Future[DummyModelRunnerOutput] = Future()

@@ -158,3 +160,40 @@ def test_async_aggregate_workers_output():

160 assert aggregated.finished_sending is None

161 assert aggregated.finished_recving == {"req2"}

162 assert aggregated.invalid_block_ids == {3, 4, 5}

163 +

164 +

165 +def test_aggregate_workers_output_with_expected_finished_count():

166 + # We create the aggregator expecting to collect from 4 workers

167 + aggregator = KVOutputAggregator(expected_finished_count=4)

168 + assert aggregator._expected_finished_count == 4

169 + # Some request with default expected finished requests

170 + output1 = DummyModelRunnerOutput(finished_sending={"req1"})

171 + aggregated = aggregator.aggregate([output1])

172 + # still expecting to collect from 4 workers

173 + assert aggregator._send_remaining_count["req1"] == 3

174 + assert not aggregated.kv_connector_output.finished_sending

175 + assert not aggregated.kv_connector_output.finished_recving

176 +

177 + # Workers discover and find that in this setup they only need to

178 + # collect from 2

179 + output1 = DummyModelRunnerOutput(

180 + finished_sending={"req1"}, expected_finished_count=2

181 + )

182 + output2 = DummyModelRunnerOutput(

183 + finished_recving={"req2"}, expected_finished_count=2

184 + )

185 + output3 = DummyModelRunnerOutput(finished_recving={"req2"})

186 + # Req2 only needs 2 acks

187 + aggregated = aggregator.aggregate([output1, output2, output3])

188 + assert aggregated.kv_connector_output.expected_finished_count == 2

189 +

190 + assert not aggregated.kv_connector_output.finished_sending

191 +

192 + # Req2 is finished

193 + assert "req2" not in aggregator._recv_remaining_count

194 + assert aggregated.kv_connector_output.finished_recving == {"req2"}

195 +

196 + # Req1 is still waiting for 2 more acks (expected_finished_count has no effect)

197 + # NOTE: This is to showcase dynamic update. Workers are responsible for

198 + # ensuring "req1" termination in this case

199 + assert aggregator._send_remaining_count["req1"] == 2

vllm/distributed/kv_transfer/kv_connector/utils.py

+30 -6

@@ -4,10 +4,9 @@

4 KV cache helper for store.

5 """

7 -from collections import defaultdict

7 from collections.abc import Sequence

8 from concurrent.futures import CancelledError, Future

10 -from typing import Literal, cast

9 +from typing import TYPE_CHECKING, Literal, cast

11 import torch

@@ -18,6 +17,9 @@ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

17 from vllm.logger import init_logger

18 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

20 +if TYPE_CHECKING:

21 + from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase

22 +

23 logger = init_logger(__name__)

@@ -124,11 +126,16 @@ class KVOutputAggregator:

126 """Utility class to aggregate the output of all workers into a single

127 output corresponding to Rank 0 for scheduler."""

128

127 - def __init__(self, world_size: int):

129 + def __init__(self, expected_finished_count: int):

130 # Complete transfer tracker. Used to track finished requests

131 # [req_id -> n_remaining_workers]

130 - self._recv_remaining_count = defaultdict[str, int](lambda: world_size)

131 - self._send_remaining_count = defaultdict[str, int](lambda: world_size)

132 + self._recv_remaining_count = dict[str, int]()

133 + self._send_remaining_count = dict[str, int]()

134 + self._expected_finished_count = expected_finished_count

135 +

136 + @classmethod

137 + def from_connector(cls, connector: "KVConnectorBase", world_size: int):

138 + return cls(connector.get_finished_count() or world_size)

139

140 def aggregate(

141 self, outputs: list[ModelRunnerOutput], output_rank: int = 0

@@ -141,7 +148,10 @@ class KVOutputAggregator:

148 finished_set: set[str],

149 ) -> None:

150 for req_id in req_ids or ():

144 - remaining_count_dict[req_id] -= 1

151 + remaining_count = remaining_count_dict.get(

152 + req_id, self._expected_finished_count

153 + )

154 + remaining_count_dict[req_id] = remaining_count - 1

155 if remaining_count_dict[req_id] == 0:

156 finished_set.add(req_id)

157 del remaining_count_dict[req_id]

@@ -154,6 +164,19 @@ class KVOutputAggregator:

164 kv_output = model_runner_output.kv_connector_output

165 if not kv_output:

166 continue

167 + # Allow the worker to dynamically update the expected number of

168 + # finished sending/recving for new requests.

169 + if (

170 + kv_output.expected_finished_count > 0

171 + and kv_output.expected_finished_count != self._expected_finished_count

172 + ):

173 + logger.debug(

174 + "Expected finished requests updated from %d to %d",

175 + self._expected_finished_count,

176 + kv_output.expected_finished_count,

177 + )

178 + self._expected_finished_count = kv_output.expected_finished_count

179 +

180 update_finished_set(

181 kv_output.finished_sending, self._send_remaining_count, finished_sending

182 )

@@ -186,6 +209,7 @@ class KVOutputAggregator:

209 finished_recving=finished_recving or None,

210 kv_connector_stats=aggregated_kv_connector_stats or None,

211 invalid_block_ids=invalid_block_ids,

212 + expected_finished_count=self._expected_finished_count,

213 )

214

215 return output

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+2 -1

@@ -413,7 +413,8 @@ class KVConnectorBase_V1(ABC):

413 def get_finished_count(self) -> int | None:

414 """

415 Get the count of requests expected to complete send/receive operations

416 - via this connector.

416 + via this connector. This method is used to initialize the

417 + KVOutputAggregator, overwriting the default world_size.

418

419 Returns:

420 int: expected sending or receiving completion count.

vllm/v1/engine/core.py

+1 -3

@@ -160,9 +160,7 @@ class EngineCore:

160 )

161 self.use_spec_decode = vllm_config.speculative_config is not None

162 if self.scheduler.connector is not None: # type: ignore

163 - self.model_executor.init_kv_output_aggregator(

164 - self.scheduler.connector.get_finished_count() # type: ignore

165 - )

163 + self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore

164

165 self.mm_registry = mm_registry = MULTIMODAL_REGISTRY

166 self.mm_receiver_cache = engine_receiver_cache_from_config(

vllm/v1/executor/abstract.py

+7 -4

@@ -5,7 +5,7 @@ from abc import ABC, abstractmethod

5 from collections.abc import Callable

6 from concurrent.futures import Future

7 from functools import cached_property

8 -from typing import Literal, TypeVar, overload

8 +from typing import TYPE_CHECKING, Literal, TypeVar, overload

10 from vllm.config import VllmConfig

11 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

@@ -19,6 +19,9 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

19 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

20 from vllm.v1.worker.worker_base import WorkerBase

22 +if TYPE_CHECKING:

23 + from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase

24 +

25 logger = init_logger(__name__)

27 _R = TypeVar("_R")

@@ -233,10 +236,10 @@ class Executor(ABC):

236 """Shutdown the executor."""

237 self.collective_rpc("shutdown")

238

236 - def init_kv_output_aggregator(self, finished_count: int | None) -> None:

239 + def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None:

240 """Init KVOutputAggregator"""

238 - self.kv_output_aggregator = KVOutputAggregator(

239 - finished_count or self.parallel_config.world_size

241 + self.kv_output_aggregator = KVOutputAggregator.from_connector(

242 + connector, self.parallel_config.world_size

243 )

244

245 @cached_property # Avoid unnecessary RPC calls

vllm/v1/outputs.py

+7 -1

@@ -86,8 +86,14 @@ class KVConnectorOutput:

86 finished_recving: set[str] | None = None

87 kv_connector_stats: KVConnectorStats | None = None

88 # IDs of externally computed KV blocks that failed to load.

89 - # Requests referencing these blocks should be rescheduled to recompute them.

89 + # Requests referencing these blocks should be rescheduled to recompute them

90 invalid_block_ids: set[int] = field(default_factory=set)

91 + # Configuration describing how many finished sending/receiving

92 + # notifications should be expected for each request. This allows

93 + # handshake-based connectors like Nixl to update the KVOutputAggregator.

94 + # It captures a static setup info and should almost always remain constant

95 + # for a given connector after discovery. Default value entails no change.

96 + expected_finished_count: int = 0

98 def is_empty(self):

99 return (