@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e53283b1c

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 2, 12:15 UTC

[Core] Cleanup KVConnector handling with PP + fix MRV2 (#43732)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nick Hill committed May 28, 2026 at 13:12 UTC 7e53283b1c28868ea01d88dd20504b61ea971bae

9 files changed +38 -48

tests/distributed/test_comm_ops.py

+1 -1

@@ -292,7 +292,7 @@ def test_async_intermediate_tensors_lazy_wait() -> None:

292 )

293

294 # accessing non-tensor attributes should not trigger wait.

295 - assert it.kv_connector_output is None

295 + assert it._comm_handles is not None

296 assert work.wait_calls == 0

297 assert post_calls["n"] == 0

298

tests/v1/worker/test_gpu_model_runner.py

+5 -2

@@ -39,6 +39,7 @@ from vllm.v1.kv_cache_interface import (

39 KVCacheGroupSpec,

40 KVCacheTensor,

41 )

42 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

43 from vllm.v1.sample.metadata import SamplingMetadata

44 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

45 from vllm.v1.worker.gpu_input_batch import InputBatch

@@ -279,7 +280,8 @@ def test_sample_tokens_receives_pp_sampled_ids_only_on_non_last_rank(

280 lambda: SimpleNamespace(world_size=world_size, is_last_rank=is_last_rank),

281 )

282

282 - assert GPUModelRunner.sample_tokens(runner, None) is None

283 + output = GPUModelRunner.sample_tokens(runner, None)

284 + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)

285 assert receive_calls == expected_calls

286

287

@@ -298,7 +300,8 @@ def test_sample_tokens_skips_pp_group_lookup_without_async_scheduling(

300 pytest.fail,

301 )

302

301 - assert GPUModelRunner.sample_tokens(runner, None) is None

303 + output = GPUModelRunner.sample_tokens(runner, None)

304 + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)

305

306

307 def test_select_common_block_size_no_valid_option():

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

+7 -2

@@ -7,6 +7,7 @@ from typing import Any

8 import torch

10 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

11 from vllm.v1.worker.gpu import eplb_utils as eplb

12 from vllm.v1.worker.gpu import model_runner as mrv2

@@ -76,7 +77,10 @@ def _make_runner(**overrides: Any) -> Any:

77 runner.max_num_reqs = 8

78 runner.max_num_tokens = 16

79 runner.decode_query_len = 1

79 - runner.kv_connector = SimpleNamespace(set_disabled=lambda *_: None)

80 + runner.kv_connector = SimpleNamespace(

81 + set_disabled=lambda *_: None,

82 + post_forward=lambda *_, **__: None,

83 + )

84 runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)

85 runner.pooling_runner = None

86 runner.execute_model_state = None

@@ -183,5 +187,6 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):

187 ),

188 )

189

186 - assert mrv2.GPUModelRunner.sample_tokens(runner, None) is None

190 + output = mrv2.GPUModelRunner.sample_tokens(runner, None)

191 + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)

192 assert events == ["postprocess", "eplb"]

vllm/sequence.py

-11

@@ -3,15 +3,9 @@

3 """Sequence and its related classes."""

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Any

7 import torch

10 -if TYPE_CHECKING:

11 - from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput

12 -else:

13 - KVConnectorOutput = Any

14 -

10 # cannot use msgspec.Struct here because Dynamo does not support it

11 @dataclass

@@ -19,24 +13,19 @@ class IntermediateTensors:

13 """For all pipeline stages except the last, we need to return the hidden

14 states and residuals to be sent to the next stage. This data structure

15 contains the hidden states and residuals for a request.

22 -

23 - Each stage also needs to handle its own kv_connector_output.

16 """

18 tensors: dict[str, torch.Tensor]

27 - kv_connector_output: KVConnectorOutput | None

20 def __init__(

21 self,

22 tensors: dict[str, torch.Tensor],

32 - kv_connector_output: KVConnectorOutput | None = None,

23 ) -> None:

24 # manually define this function, so that

25 # Dynamo knows `IntermediateTensors()` comes from this file.

26 # Otherwise, dataclass will generate this function by evaluating

27 # a string, and we will lose the information about the source file.

28 self.tensors = tensors

39 - self.kv_connector_output = kv_connector_output

30 def __getitem__(self, key: str | slice):

31 if isinstance(key, str):

vllm/v1/outputs.py

+14

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from abc import ABC, abstractmethod

5 +from copy import copy

6 from dataclasses import dataclass, field

7 from typing import TYPE_CHECKING, NamedTuple, TypeAlias

@@ -279,6 +280,19 @@ class ModelRunnerOutput:

280 # ``None`` when ``enable_return_routed_experts`` is off.

281 routed_experts: RoutedExpertsLists | None = None

282

283 + @staticmethod

284 + def with_kv_conn_output_only(

285 + kv_connector_output: KVConnectorOutput | None,

286 + ) -> "ModelRunnerOutput":

287 + """Return ModelRunnerOutput containing the provided KVConnectorOutput,

288 + otherwise empty. Returns None if kv_connector_output is passed as None.

289 + """

290 + if kv_connector_output is None or kv_connector_output.is_empty():

291 + return EMPTY_MODEL_RUNNER_OUTPUT

292 + output = copy(EMPTY_MODEL_RUNNER_OUTPUT)

293 + output.kv_connector_output = kv_connector_output

294 + return output

295 +

296

297 # ModelRunnerOutput wrapper for async scheduling.

298 class AsyncModelRunnerOutput(ABC):

vllm/v1/worker/gpu/kv_connector.py

+1 -6

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import copy

3 from typing import TYPE_CHECKING

5 import torch

@@ -103,11 +102,7 @@ class ActiveKVConnector(KVConnector):

102 self.pre_forward(scheduler_output)

103 finished_req_ids = scheduler_output.finished_req_ids

104 kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)

106 - if kv_connector_output is None or kv_connector_output.is_empty():

107 - return EMPTY_MODEL_RUNNER_OUTPUT

108 - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

109 - output.kv_connector_output = kv_connector_output

110 - return output

105 + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

106

107 def set_disabled(self, disabled: bool) -> None:

108 # Ensure that layer-wise connector hooks aren't called when disabled.

vllm/v1/worker/gpu/model_runner.py

+8 -8

@@ -1218,9 +1218,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1218

1219 if not self.is_last_pp_rank:

1220 # Non-last PP rank: return IntermediateTensors for sending.

1221 - assert output_intermediate_tensors is not None

1222 - kv_connector_output = self.kv_connector.post_forward(finished_req_ids)

1223 - output_intermediate_tensors.kv_connector_output = kv_connector_output

1221 return output_intermediate_tensors

1222 return None

1223

@@ -1249,7 +1246,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1246 input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1

1247 )

1248 self.postprocess(input_batch, sampled, num_sampled, num_rejected)

1252 - return None

1249 +

1250 + # Post-step KV connector related operations.

1251 + kv_connector_output = self.kv_connector.post_forward(finished_req_ids)

1252 + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

1253

1254 # Last rank: sample tokens

1255 sampler_output, num_sampled, num_rejected = self.sample(

@@ -1367,18 +1367,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1367 finished_req_ids = self.execute_model_state.finished_req_ids

1368 self.execute_model_state = None

1369

1370 + # Post-step KV connector related operations.

1371 + kv_connector_output = self.kv_connector.post_forward(finished_req_ids)

1372 +

1373 if not self.is_last_pp_rank:

1374 self.postprocess_pool(input_batch)

1372 - return None

1375 + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

1376

1377 assert self.pooling_runner is not None

1378 pooler_output, is_valid = self.pooling_runner.pool(

1379 hidden_states, input_batch, self.req_states

1380 )

1381

1379 - # Post-step KV connector related operations.

1380 - kv_connector_output = self.kv_connector.post_forward(finished_req_ids)

1381 -

1382 # Build the model runner output.

1383 model_runner_output = ModelRunnerOutput(

1384 req_ids=input_batch.req_ids,

vllm/v1/worker/gpu_model_runner.py

+1 -10

@@ -4250,7 +4250,6 @@ class GPUModelRunner(

4250 if not get_pp_group().is_last_rank:

4251 # Return the intermediate tensors.

4252 assert isinstance(hidden_states, IntermediateTensors)

4253 - hidden_states.kv_connector_output = kv_connector_output

4253 self.kv_connector_output = kv_connector_output

4254 return hidden_states

4255

@@ -4326,17 +4325,9 @@ class GPUModelRunner(

4325 # receive sampled token ids from the last PP rank.

4326 if self.use_async_scheduling and not get_pp_group().is_last_rank:

4327 self._pp_receive_prev_sampled_token_ids_to_input_batch()

4329 - if not kv_connector_output:

4330 - return None # type: ignore[return-value]

4331 -

4328 # In case of PP with kv transfer, we need to pass through the

4329 # kv_connector_output

4334 - if kv_connector_output.is_empty():

4335 - return EMPTY_MODEL_RUNNER_OUTPUT

4336 -

4337 - output = copy(EMPTY_MODEL_RUNNER_OUTPUT)

4338 - output.kv_connector_output = kv_connector_output

4339 - return output

4330 + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

4331

4332 # Unpack ephemeral state.

4333 (

vllm/v1/worker/kv_connector_model_runner_mixin.py

+1 -8

@@ -4,7 +4,6 @@

4 Define KV connector functionality mixin for model runners.

5 """

7 -import copy

7 from collections.abc import Generator

8 from contextlib import AbstractContextManager, contextmanager, nullcontext

9 from typing import TYPE_CHECKING

@@ -20,7 +19,6 @@ from vllm.logger import init_logger

19 from vllm.v1.attention.backend import AttentionBackend

20 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig

21 from vllm.v1.outputs import (

23 - EMPTY_MODEL_RUNNER_OUTPUT,

22 KVConnectorOutput,

23 ModelRunnerOutput,

24 )

@@ -47,12 +45,7 @@ class KVConnectorModelRunnerMixin:

45 ):

46 pass

50 - if kv_connector_output.is_empty():

51 - return EMPTY_MODEL_RUNNER_OUTPUT

52 -

53 - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

54 - output.kv_connector_output = kv_connector_output

55 - return output

48 + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

50 @staticmethod

51 def maybe_get_kv_connector_output(