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

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

[V1] [P/D] Refactor KV Connector Path (#21980)

Signed-off-by: David Ben-David <davidb@pliops.com> Co-authored-by: David Ben-David <davidb@pliops.com>

David Ben-David committed Aug 3, 2025 at 14:03 UTC aefeea0fde0fbe5871a0799fad583e6ed6fdf903

12 files changed +141 -79

tests/v1/kv_connector/unit/test_output_aggreagator.py

+17 -3

@@ -4,7 +4,7 @@ from concurrent.futures import Future

4 from typing import Optional

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

7 -from vllm.v1.outputs import ModelRunnerOutput

7 +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

10 class DummyModelRunnerOutput(ModelRunnerOutput):

@@ -12,8 +12,16 @@ class DummyModelRunnerOutput(ModelRunnerOutput):

12 def __init__(self,

13 finished_sending: Optional[set[str]] = None,

14 finished_recving: Optional[set[str]] = None):

15 - self.finished_sending = finished_sending

16 - self.finished_recving = finished_recving

15 + self.kv_connector_output = KVConnectorOutput(

16 + finished_sending=finished_sending,

17 + finished_recving=finished_recving,

18 + )

19 +

20 + def __repr__(self):

21 + return (

22 + f"DummyModelRunnerOutput("

23 + f"finished_sending={self.kv_connector_output.finished_sending},"

24 + f"finished_recving={self.kv_connector_output.finished_recving})")

27 def test_aggregate_workers_output():

@@ -27,6 +35,7 @@ def test_aggregate_workers_output():

35 aggregated = aggregator.aggregate([output1, output2])

37 assert aggregated is output1

38 + aggregated = aggregated.kv_connector_output

39 assert aggregated.finished_sending is None

40 assert aggregated.finished_recving is None

@@ -38,6 +47,7 @@ def test_aggregate_workers_output():

47 aggregated = aggregator.aggregate([output1, output2])

49 assert aggregated is output1

50 + aggregated = aggregated.kv_connector_output

51 assert aggregated.finished_sending == {'req1'}

52 assert aggregated.finished_recving is None

@@ -49,6 +59,7 @@ def test_aggregate_workers_output():

59 aggregated = aggregator.aggregate([output1, output2])

61 assert aggregated is output1

62 + aggregated = aggregated.kv_connector_output

63 assert aggregated.finished_sending is None

64 assert aggregated.finished_recving == {'req2'}

@@ -70,6 +81,7 @@ def test_async_aggregate_workers_output():

81 assert result_future.done()

82 aggregated = result_future.result()

83 assert aggregated is output1

84 + aggregated = aggregated.kv_connector_output

85 assert aggregated.finished_sending is None

86 assert aggregated.finished_recving is None

@@ -87,6 +99,7 @@ def test_async_aggregate_workers_output():

99 assert result_future.done()

100 aggregated = result_future.result()

101 assert aggregated is output1

102 + aggregated = aggregated.kv_connector_output

103 assert aggregated.finished_sending == {'req1'}

104 assert aggregated.finished_recving is None

105

@@ -104,5 +117,6 @@ def test_async_aggregate_workers_output():

117 assert result_future.done()

118 aggregated = result_future.result()

119 assert aggregated is output1

120 + aggregated = aggregated.kv_connector_output

121 assert aggregated.finished_sending is None

122 assert aggregated.finished_recving == {'req2'}

tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py

+5 -3

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

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

3 import copy

5 -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

5 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput

6 from vllm.v1.request import FinishReason, RequestStatus

8 from .utils import (assert_scheduler_empty, create_model_runner_output,

@@ -86,7 +86,8 @@ def test_basic_lifecycle():

87 # (3b): execute_model()

88 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)

89 - model_runner_output.finished_sending = [request_id]

89 + model_runner_output.kv_connector_output = KVConnectorOutput(

90 + finished_sending=[request_id])

92 # (3c): update_from_output()

93 scheduler.update_from_output(scheduler_output, model_runner_output)

@@ -176,7 +177,8 @@ def test_prefix_cache_lifecycle():

177 scheduler_output = scheduler.schedule()

178 scheduler.schedule()

179 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)

179 - model_runner_output.finished_sending = [request_remote.request_id]

180 + model_runner_output.kv_connector_output = KVConnectorOutput(

181 + finished_sending=[request_remote.request_id])

182 scheduler.update_from_output(scheduler_output, model_runner_output)

183 _ = scheduler.schedule()

184 assert_scheduler_empty(scheduler)

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+5 -3

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

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

3 import copy

5 -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

5 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput

6 from vllm.v1.request import FinishReason, RequestStatus

8 from .utils import (assert_scheduler_empty, create_model_runner_output,

@@ -72,7 +72,8 @@ def test_basic_lifecycle():

73 # (2b): forward(): request finishes recv.

74 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)

75 - model_runner_output.finished_recving = [request_id]

75 + model_runner_output.kv_connector_output = KVConnectorOutput(

76 + finished_recving=[request_id])

78 # (2c): update_from_output():

79 engine_core_outputs = scheduler.update_from_output(scheduler_output,

@@ -309,7 +310,8 @@ def test_full_block_prompt():

310 # # STEP (2): Recv.

311 scheduler_output = scheduler.schedule()

312 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)

312 - model_runner_output.finished_recving = [request_id]

313 + model_runner_output.kv_connector_output = KVConnectorOutput(

314 + finished_recving=[request_id])

315 scheduler.update_from_output(scheduler_output, model_runner_output)

316 assert len(scheduler.waiting) == 1

317 assert (request_id in scheduler.finished_recving_kv_req_ids)

tests/v1/kv_connector/unit/utils.py

+5 -3

@@ -17,7 +17,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheBlocks

17 from vllm.v1.core.sched.scheduler import Scheduler

18 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

19 KVCacheGroupSpec)

20 -from vllm.v1.outputs import ModelRunnerOutput

20 +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

21 from vllm.v1.request import Request

22 from vllm.v1.structured_output import StructuredOutputManager

@@ -188,8 +188,10 @@ def create_model_runner_output(

188 logprobs=None,

189 prompt_logprobs_dict={},

190 pooler_output=None,

191 - finished_sending=finished_sending,

192 - finished_recving=finished_recving,

191 + kv_connector_output=KVConnectorOutput(

192 + finished_sending=finished_sending,

193 + finished_recving=finished_recving,

194 + ),

195 )

196

197

vllm/distributed/kv_transfer/kv_connector/utils.py

+7 -9

@@ -16,7 +16,7 @@ from vllm.config import VllmConfig, get_current_vllm_config

16 from vllm.distributed.kv_transfer.kv_connector.factory import (

17 KVConnectorFactory)

18 from vllm.logger import init_logger

19 -from vllm.v1.outputs import ModelRunnerOutput

19 +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

21 logger = init_logger(__name__)

@@ -129,7 +129,7 @@ class KVOutputAggregator:

129 def aggregate(self,

130 outputs: list[ModelRunnerOutput],

131 output_rank: int = 0) -> ModelRunnerOutput:

132 - # aggregate finished_sending, finished_recving from all workers

132 + # aggregate kv_connector_output from all workers

133

134 def update_finished_set(req_ids: Optional[set[str]],

135 remaining_count_dict: dict[str, int],

@@ -143,6 +143,7 @@ class KVOutputAggregator:

143 finished_sending = set[str]()

144 finished_recving = set[str]()

145 for output in outputs:

146 + output = output.kv_connector_output

147 update_finished_set(output.finished_sending,

148 self._send_remaining_count, finished_sending)

149 update_finished_set(output.finished_recving,

@@ -151,13 +152,10 @@ class KVOutputAggregator:

152 # select output of the worker specified by output_rank

153 output = outputs[output_rank]

154

154 - # set the aggregated finished_sending / finished_recving

155 - # if output.finished_sending/recving is not empty, but the other ranks

156 - # still have unfinished send/recv, we want to set the aggregated

157 - # finished_sending/recving to None until all ranks have finished

158 - # send/recv

159 - output.finished_sending = finished_sending if finished_sending else None

160 - output.finished_recving = finished_recving if finished_recving else None

155 + output.kv_connector_output = KVConnectorOutput(

156 + finished_sending=finished_sending or None,

157 + finished_recving=finished_recving or None,

158 + )

159

160 return output

161

vllm/sequence.py

+7 -6

@@ -10,7 +10,7 @@ from collections.abc import Mapping

10 from collections.abc import Sequence as GenericSequence

11 from dataclasses import dataclass, field

12 from functools import reduce

13 -from typing import Any, Callable, Optional, Union

13 +from typing import TYPE_CHECKING, Any, Callable, Optional, Union

15 import msgspec

16 import torch

@@ -21,6 +21,10 @@ from vllm.multimodal import MultiModalKwargs, MultiModalPlaceholderDict

21 from vllm.pooling_params import PoolingParams

22 from vllm.sampling_params import RequestOutputKind, SamplingParams

24 +if TYPE_CHECKING:

25 + from vllm.v1.worker.kv_connector_model_runner_mixin import (

26 + KVConnectorOutput)

27 +

28 VLLM_TOKEN_ID_ARRAY_TYPE = "l"

30 VLLM_INVALID_TOKEN_ID = -1

@@ -1159,14 +1163,11 @@ class IntermediateTensors:

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

1164 contains the hidden states and residuals for a request.

1165

1162 - Each stage also needs to handle its own finished_sending and

1163 - finished_recving in case of kv transfer.

1166 + Each stage also needs to handle its own kv_connector_output.

1167 """

1168

1169 tensors: dict[str, torch.Tensor]

1167 - # [req_ids]

1168 - finished_sending: Optional[set[str]] = None

1169 - finished_recving: Optional[set[str]] = None

1170 + kv_connector_output: Optional["KVConnectorOutput"]

1171

1172 def __init__(self, tensors):

1173 # manually define this function, so that

vllm/v1/core/sched/scheduler.py

+7 -5

@@ -30,7 +30,7 @@ from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,

30 EngineCoreOutputs)

31 from vllm.v1.kv_cache_interface import KVCacheConfig

32 from vllm.v1.metrics.stats import SchedulerStats

33 -from vllm.v1.outputs import ModelRunnerOutput

33 +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

34 from vllm.v1.request import Request, RequestStatus

35 from vllm.v1.spec_decode.metrics import SpecDecodingStats

36 from vllm.v1.structured_output import StructuredOutputManager

@@ -884,7 +884,9 @@ class Scheduler(SchedulerInterface):

884 self.waiting.remove_requests(stopped_preempted_reqs)

885

886 # KV Connector: update state for finished KV Transfers.

887 - self._update_from_kv_xfer_finished(model_runner_output)

887 + if model_runner_output.kv_connector_output:

888 + self._update_from_kv_xfer_finished(

889 + model_runner_output.kv_connector_output)

890

891 # Create EngineCoreOutputs for all clients that have requests with

892 # outputs in this step.

@@ -1128,7 +1130,7 @@ class Scheduler(SchedulerInterface):

1130 return True

1131

1132 def _update_from_kv_xfer_finished(self,

1131 - model_runner_output: ModelRunnerOutput):

1133 + kv_connector_output: KVConnectorOutput):

1134 """

1135 KV Connector: update the scheduler state based on the output.

1136

@@ -1139,9 +1141,9 @@ class Scheduler(SchedulerInterface):

1141 scheduler the request during the next step.

1142 """

1143 # KV Connector:: update recv and send status from last step.

1142 - for req_id in (model_runner_output.finished_recving or ()):

1144 + for req_id in (kv_connector_output.finished_recving or ()):

1145 logger.debug("Finished recving KV transfer for request %s", req_id)

1146 self.finished_recving_kv_req_ids.add(req_id)

1145 - for req_id in (model_runner_output.finished_sending or ()):

1147 + for req_id in (kv_connector_output.finished_sending or ()):

1148 logger.debug("Finished sending KV transfer for request %s", req_id)

1149 self._free_blocks(self.requests[req_id])

vllm/v1/outputs.py

+8 -5

@@ -71,6 +71,13 @@ class SamplerOutput:

71 logprobs_tensors: Optional[LogprobsTensors]

74 +@dataclass

75 +class KVConnectorOutput:

76 + # [req_ids]

77 + finished_sending: Optional[set[str]] = None

78 + finished_recving: Optional[set[str]] = None

79 +

80 +

81 # ModelRunnerOutput is serialized and sent to the scheduler process.

82 # This is expensive for torch.Tensor so prefer to use list instead.

83 @dataclass

@@ -104,9 +111,7 @@ class ModelRunnerOutput:

111 # [num_reqs, hidden_size]

112 pooler_output: list[Optional[torch.Tensor]]

113

107 - # [req_ids]

108 - finished_sending: Optional[set[str]] = None

109 - finished_recving: Optional[set[str]] = None

114 + kv_connector_output: Optional[KVConnectorOutput] = None

115

116 # req_id -> num_nans_in_logits

117 num_nans_in_logits: Optional[dict[str, int]] = None

@@ -119,6 +124,4 @@ EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[],

124 logprobs=None,

125 prompt_logprobs_dict={},

126 pooler_output=[],

122 - finished_sending=None,

123 - finished_recving=None,

127 num_nans_in_logits=None)

vllm/v1/worker/gpu_model_runner.py

+9 -19

@@ -69,7 +69,7 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

69 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

70 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

71 from vllm.v1.worker.kv_connector_model_runner_mixin import (

72 - KVConnectorModelRunnerMixin)

72 + KVConnectorModelRunnerMixin, KVConnectorOutput)

73 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

75 from ..sample.logits_processor import LogitsProcessorManager

@@ -1423,8 +1423,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1423 hidden_states: torch.Tensor,

1424 num_scheduled_tokens: int,

1425 num_scheduled_tokens_np: np.ndarray,

1426 - finished_sending: Optional[set[str]],

1427 - finished_recving: Optional[set[str]],

1426 + kv_connector_output: Optional[KVConnectorOutput],

1427 ) -> ModelRunnerOutput:

1428 assert self.input_batch.num_reqs ==\

1429 len(self.input_batch.pooling_params), \

@@ -1459,8 +1458,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1458 logprobs=None,

1459 prompt_logprobs_dict={},

1460 pooler_output=pooler_output,

1462 - finished_sending=finished_sending,

1463 - finished_recving=finished_recving,

1461 + kv_connector_output=kv_connector_output,

1462 )

1463

1464 @torch.inference_mode()

@@ -1564,8 +1562,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1562 num_tokens=num_input_tokens,

1563 num_tokens_across_dp=num_tokens_across_dp,

1564 skip_cuda_graphs=skip_cuda_graphs,

1567 - ):

1568 - self.maybe_setup_kv_connector(scheduler_output)

1565 + ), self.maybe_get_kv_connector_output(

1566 + scheduler_output) as kv_connector_output:

1567

1568 model_output = self.model(

1569 input_ids=input_ids,

@@ -1578,10 +1576,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1576 ),

1577 )

1578

1581 - self.maybe_wait_for_kv_save()

1582 - finished_sending, finished_recving = (

1583 - self.get_finished_kv_transfers(scheduler_output))

1584 -

1579 if self.use_aux_hidden_state_outputs:

1580 hidden_states, aux_hidden_states = model_output

1581 else:

@@ -1597,20 +1591,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1591 == "external_launcher" and len(get_pp_group().ranks) > 0

1592 if not get_pp_group().is_last_rank:

1593 # For mid-pipeline stages, return the hidden states.

1594 + assert isinstance(hidden_states, IntermediateTensors)

1595 if not broadcast_pp_output:

1601 - if finished_sending or finished_recving:

1602 - hidden_states.finished_sending = finished_sending

1603 - hidden_states.finished_recving = finished_recving

1596 + hidden_states.kv_connector_output = kv_connector_output

1597 return hidden_states

1605 - assert isinstance(hidden_states, IntermediateTensors)

1598 get_pp_group().send_tensor_dict(hidden_states.tensors,

1599 all_gather_group=get_tp_group())

1600 logits = None

1601 else:

1602 if self.input_batch.pooling_params:

1603 return self._pool(hidden_states, num_scheduled_tokens,

1612 - num_scheduled_tokens_np, finished_sending,

1613 - finished_recving)

1604 + num_scheduled_tokens_np, kv_connector_output)

1605

1606 sample_hidden_states = hidden_states[logits_indices]

1607 logits = self.model.compute_logits(sample_hidden_states, None)

@@ -1760,8 +1751,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1751 logprobs=logprobs_lists,

1752 prompt_logprobs_dict=prompt_logprobs_dict,

1753 pooler_output=[],

1763 - finished_sending=finished_sending,

1764 - finished_recving=finished_recving,

1754 + kv_connector_output=kv_connector_output,

1755 num_nans_in_logits=num_nans_in_logits,

1756 )

1757

vllm/v1/worker/gpu_worker.py

+12 -10

@@ -16,8 +16,7 @@ from vllm.config import VllmConfig

16 from vllm.distributed import (ensure_model_parallel_initialized,

17 init_distributed_environment,

18 set_custom_all_reduce)

19 -from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,

20 - has_kv_transfer_group)

19 +from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized

20 from vllm.distributed.parallel_state import get_pp_group, get_tp_group

21 from vllm.logger import init_logger

22 from vllm.lora.request import LoRARequest

@@ -369,17 +368,20 @@ class Worker(WorkerBase):

368 assert isinstance(output, IntermediateTensors)

369 get_pp_group().send_tensor_dict(output.tensors,

370 all_gather_group=get_tp_group())

372 - if not has_kv_transfer_group():

371 +

372 + kv_connector_output = output.kv_connector_output

373 + if not kv_connector_output:

374 return None

375

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

376 - # finished_sending and finished_recving buffers.

377 - new_output = EMPTY_MODEL_RUNNER_OUTPUT

378 - if output.finished_sending or output.finished_recving:

379 - new_output = copy.copy(new_output)

380 - new_output.finished_sending = output.finished_sending

381 - new_output.finished_recving = output.finished_recving

382 - output = new_output

377 + # kv_connector_output

378 + if (not kv_connector_output.finished_sending

379 + and not kv_connector_output.finished_recving):

380 + return EMPTY_MODEL_RUNNER_OUTPUT

381 +

382 + output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

383 + output.kv_connector_output = kv_connector_output

384 + return output

385

386 assert isinstance(output, ModelRunnerOutput)

387 return output

vllm/v1/worker/kv_connector_model_runner_mixin.py

+54 -9

@@ -4,6 +4,8 @@

4 Define KV connector functionality mixin for model runners.

5 """

6 import copy

7 +from contextlib import AbstractContextManager, contextmanager, nullcontext

8 +from typing import Generator # noqa: UP035

9 from typing import TYPE_CHECKING, Optional

11 from vllm.config import VllmConfig

@@ -12,7 +14,8 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,

14 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

15 from vllm.forward_context import get_forward_context, set_forward_context

16 from vllm.logger import init_logger

15 -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput

17 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput,

18 + ModelRunnerOutput)

20 if TYPE_CHECKING:

21 from vllm.v1.core.sched.output import SchedulerOutput

@@ -53,18 +56,60 @@ class KVConnectorModelRunnerMixin:

56 scheduler_output.finished_req_ids)

57 return None, None

56 - def kv_connector_no_forward(self, scheduler_output: "SchedulerOutput",

59 + @staticmethod

60 + def kv_connector_no_forward(scheduler_output: "SchedulerOutput",

61 vllm_config: VllmConfig) -> ModelRunnerOutput:

62 # KV send/recv even if no work to do.

59 - with set_forward_context(None, vllm_config):

60 - self.maybe_setup_kv_connector(scheduler_output)

61 - finished_sending, finished_recving = (

62 - self.get_finished_kv_transfers(scheduler_output))

63 + with set_forward_context(

64 + None, vllm_config

65 + ), KVConnectorModelRunnerMixin._get_kv_connector_output(

66 + scheduler_output, wait_for_save=False) as kv_connector_output:

67 + pass

64 - if not finished_sending and not finished_recving:

69 + if (not kv_connector_output.finished_sending

70 + and not kv_connector_output.finished_recving):

71 return EMPTY_MODEL_RUNNER_OUTPUT

73 output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

68 - output.finished_sending = finished_sending

69 - output.finished_recving = finished_recving

74 + output.kv_connector_output = kv_connector_output

75 return output

76 +

77 + @staticmethod

78 + def maybe_get_kv_connector_output(

79 + scheduler_output: "SchedulerOutput"

80 + ) -> AbstractContextManager[Optional[KVConnectorOutput]]:

81 + return KVConnectorModelRunnerMixin._get_kv_connector_output(

82 + scheduler_output) if has_kv_transfer_group() else nullcontext()

83 +

84 + # This context manager must be used within an active forward context.

85 + # It encapsulates the entire KV conector lifecycle within execute_model

86 + @staticmethod

87 + @contextmanager

88 + def _get_kv_connector_output(

89 + scheduler_output: "SchedulerOutput",

90 + wait_for_save: bool = True

91 + ) -> Generator[KVConnectorOutput, None, None]:

92 + output = KVConnectorOutput()

93 +

94 + # Update KVConnector with the KVConnector metadata forward().

95 + kv_connector = get_kv_transfer_group()

96 + assert isinstance(kv_connector, KVConnectorBase_V1)

97 + assert scheduler_output.kv_connector_metadata is not None

98 + kv_connector.bind_connector_metadata(

99 + scheduler_output.kv_connector_metadata)

100 +

101 + # Background KV cache transfers happen here.

102 + # These transfers are designed to be async and the requests

103 + # involved may be disjoint from the running requests.

104 + # Do this here to save a collective_rpc.

105 + kv_connector.start_load_kv(get_forward_context())

106 + try:

107 + yield output

108 + finally:

109 + if wait_for_save:

110 + kv_connector.wait_for_save()

111 +

112 + output.finished_sending, output.finished_recving = (

113 + kv_connector.get_finished(scheduler_output.finished_req_ids))

114 +

115 + kv_connector.clear_connector_metadata()

vllm/v1/worker/tpu_model_runner.py

+5 -4

@@ -51,7 +51,7 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsLists,

51 from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata

52 from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler

53 from vllm.v1.worker.kv_connector_model_runner_mixin import (

54 - KVConnectorModelRunnerMixin)

54 + KVConnectorModelRunnerMixin, KVConnectorOutput)

55 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

56 from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch

@@ -1175,9 +1175,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1175 logprobs=logprobs_lists,

1176 prompt_logprobs_dict=prompt_logprobs_dict,

1177 pooler_output=[],

1178 - finished_sending=finished_sending,

1179 - finished_recving=finished_recving,

1180 - )

1178 + kv_connector_output=KVConnectorOutput(

1179 + finished_sending=finished_sending,

1180 + finished_recving=finished_recving,

1181 + ))

1182

1183 # Check there are no new graphs compiled - all the graphs should be

1184 # captured and compiled during warm up.