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 9, 18:15 UTC
[Model Runner V2] Add KV Connector support (#32742)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 23, 2026 at 10:49 UTC
8518b3044794d11b6a9881f2875cfa5d2823ac56
4 files changed
+159
-16
vllm/v1/worker/gpu/attn_utils.py
+2
-1
@@ -133,10 +133,11 @@ def init_kv_cache(
133
kv_cache_config: KVCacheConfig,
134
attn_backends: dict[str, AttentionBackend],
135
device: torch.device,
136
-) -> None:
136
+) -> dict[str, torch.Tensor]:
137
kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
138
kv_caches = _reshape_kv_cache(kv_cache_config, kv_cache_raw_tensors, attn_backends)
139
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
140
+ return kv_caches
141
142
143
def build_attn_metadata(
vllm/v1/worker/gpu/kv_connector.py
new
+125
@@ -0,0 +1,125 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
+from typing import TYPE_CHECKING
5
+
6
+import torch
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.distributed.kv_transfer import (
10
+ get_kv_transfer_group,
11
+ has_kv_transfer_group,
12
+ kv_transfer_state,
13
+)
14
+from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
15
+from vllm.forward_context import (
16
+ get_forward_context,
17
+ is_forward_context_available,
18
+ set_forward_context,
19
+)
20
+from vllm.v1.outputs import (
21
+ EMPTY_MODEL_RUNNER_OUTPUT,
22
+ KVConnectorOutput,
23
+ ModelRunnerOutput,
24
+)
25
+
26
+if TYPE_CHECKING:
27
+ from vllm.v1.core.sched.output import SchedulerOutput
28
+
29
+
30
+class KVConnector:
31
+ """KVConnector interface used by GPUModelRunner."""
32
+
33
+ def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
34
+ pass
35
+
36
+ def post_forward(
37
+ self, scheduler_output: "SchedulerOutput", wait_for_save: bool = True
38
+ ) -> KVConnectorOutput | None:
39
+ return None
40
+
41
+ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
42
+ return EMPTY_MODEL_RUNNER_OUTPUT
43
+
44
+ def set_disabled(self, disabled: bool) -> None:
45
+ pass
46
+
47
+
48
+class ActiveKVConnector(KVConnector):
49
+ def __init__(
50
+ self, vllm_config: VllmConfig, kv_caches_dict: dict[str, torch.Tensor]
51
+ ):
52
+ self.vllm_config = vllm_config
53
+ self.kv_connector = get_kv_transfer_group()
54
+ # Register kv caches with KV Connector if applicable.
55
+ # TODO: support cross_layers_kv_cache
56
+ # (see https://github.com/vllm-project/vllm/pull/27743)
57
+ self.kv_connector.register_kv_caches(kv_caches_dict)
58
+ self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
59
+
60
+ self._disabled = False
61
+
62
+ def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
63
+ if self._disabled:
64
+ return
65
+
66
+ if scheduler_output.preempted_req_ids:
67
+ self.kv_connector.handle_preemptions(scheduler_output.preempted_req_ids)
68
+ assert scheduler_output.kv_connector_metadata is not None
69
+ self.kv_connector.bind_connector_metadata(
70
+ scheduler_output.kv_connector_metadata
71
+ )
72
+ # TODO: sort out KV Connectors' use of forward_context
73
+ if is_forward_context_available():
74
+ self.kv_connector.start_load_kv(get_forward_context())
75
+ else:
76
+ with set_forward_context(None, self.vllm_config):
77
+ self.kv_connector.start_load_kv(get_forward_context())
78
+
79
+ def post_forward(
80
+ self, scheduler_output: "SchedulerOutput", wait_for_save: bool = True
81
+ ) -> KVConnectorOutput | None:
82
+ if self._disabled:
83
+ return None
84
+
85
+ output = KVConnectorOutput()
86
+ if wait_for_save:
87
+ self.kv_connector.wait_for_save()
88
+ output.finished_sending, output.finished_recving = (
89
+ self.kv_connector.get_finished(scheduler_output.finished_req_ids)
90
+ )
91
+ output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
92
+ output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
93
+ output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events()
94
+ self.kv_connector.clear_connector_metadata()
95
+ return output
96
+
97
+ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
98
+ if self._disabled:
99
+ return EMPTY_MODEL_RUNNER_OUTPUT
100
+
101
+ self.pre_forward(scheduler_output)
102
+ kv_connector_output = self.post_forward(scheduler_output, wait_for_save=False)
103
+ if kv_connector_output is None or kv_connector_output.is_empty():
104
+ return EMPTY_MODEL_RUNNER_OUTPUT
105
+ output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
106
+ output.kv_connector_output = kv_connector_output
107
+ return output
108
+
109
+ def set_disabled(self, disabled: bool) -> None:
110
+ # Ensure that layer-wise connector hooks aren't called when disabled.
111
+ kv_transfer_state._KV_CONNECTOR_AGENT = None if disabled else self.kv_connector
112
+ self._disabled = disabled
113
+
114
+
115
+NO_OP_KV_CONNECTOR = KVConnector()
116
+
117
+
118
+def get_kv_connector(
119
+ vllm_config: VllmConfig, kv_caches_dict: dict[str, torch.Tensor]
120
+) -> KVConnector:
121
+ if not has_kv_transfer_group():
122
+ # No-op connector.
123
+ return NO_OP_KV_CONNECTOR
124
+
125
+ return ActiveKVConnector(vllm_config, kv_caches_dict)
vllm/v1/worker/gpu/model_runner.py
+28
-13
@@ -20,10 +20,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
20
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
21
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22
from vllm.v1.kv_cache_interface import KVCacheConfig
23
-from vllm.v1.outputs import (
24
- EMPTY_MODEL_RUNNER_OUTPUT,
25
- ModelRunnerOutput,
26
-)
23
+from vllm.v1.outputs import ModelRunnerOutput
24
from vllm.v1.worker.gpu.async_utils import AsyncOutput
25
from vllm.v1.worker.gpu.attn_utils import (
26
build_attn_metadata,
@@ -48,6 +45,11 @@ from vllm.v1.worker.gpu.input_batch import (
45
prepare_pos_seq_lens,
46
prepare_prefill_inputs,
47
)
48
+from vllm.v1.worker.gpu.kv_connector import (
49
+ NO_OP_KV_CONNECTOR,
50
+ KVConnector,
51
+ get_kv_connector,
52
+)
53
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
54
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
55
from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -57,13 +59,12 @@ from vllm.v1.worker.gpu.spec_decode import init_speculator
59
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
60
from vllm.v1.worker.gpu.states import RequestState
61
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
60
-from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
62
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
63
64
logger = init_logger(__name__)
65
66
66
-class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
67
+class GPUModelRunner(LoRAModelRunnerMixin):
68
def __init__(
69
self,
70
vllm_config: VllmConfig,
@@ -172,6 +173,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
173
self.tmp_cu_num_logits = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
174
self.tmp_query_start_loc = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
175
176
+ self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
177
+
178
def update_max_model_len(self, max_model_len: int) -> None:
179
self.max_model_len = max_model_len
180
self.req_states.max_model_len = max_model_len
@@ -248,13 +251,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
251
)
252
253
self.kv_caches: list[torch.Tensor] = []
251
- init_kv_cache(
254
+ kv_caches_dict = init_kv_cache(
255
self.kv_caches,
256
self.compilation_config.static_forward_context,
257
self.kv_cache_config,
258
self.attn_backends,
259
self.device,
260
)
261
+ self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
262
+
263
# Attention groups are not supported.
264
self.attn_groups = [] # type: ignore
265
@@ -291,18 +296,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
296
num_tokens_per_request[-1] += num_tokens % num_reqs
297
assert sum(num_tokens_per_request) == num_tokens
298
num_scheduled_tokens = {
294
- f"_dummy_req_{i}": num_tokens_per_request[i] for i in range(num_reqs)
299
+ f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
300
}
301
dummy_scheduler_output = SchedulerOutput.make_empty()
302
dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
303
dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
304
305
+ # Disable any use of KVConnector for dummy runs.
306
+ self.kv_connector.set_disabled(True)
307
+
308
# Execute the model.
309
self.execute_model(
310
dummy_scheduler_output, dummy_run=True, skip_attn_for_dummy_run=skip_attn
311
)
312
+ self.kv_connector.set_disabled(False)
313
assert self.execute_model_state is not None
305
- hidden_states, input_batch = self.execute_model_state
314
+ hidden_states, input_batch, _ = self.execute_model_state
315
sample_hidden_states = hidden_states[input_batch.logits_indices]
316
return hidden_states, sample_hidden_states
317
@@ -792,7 +801,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
801
self.block_tables.apply_staged_writes()
802
if scheduler_output.total_num_scheduled_tokens == 0:
803
# No need to run the model.
795
- return EMPTY_MODEL_RUNNER_OUTPUT
804
+ empty_output = self.kv_connector.no_forward(scheduler_output)
805
+ return empty_output
806
807
# Get the CUDA graph size. None means no CUDA graph is used.
808
cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
@@ -809,7 +819,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
819
)
820
if num_tokens_after_padding == 0:
821
# All DP ranks have zero tokens to run.
812
- return EMPTY_MODEL_RUNNER_OUTPUT
822
+ empty_output = self.kv_connector.no_forward(scheduler_output)
823
+ return empty_output
824
825
if not dummy_run:
826
# Common case.
@@ -860,6 +871,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
871
# Run CUDA graph.
872
# NOTE(woosuk): Here, we don't need to pass the input tensors,
873
# because they are already copied to the CUDA graph input buffers.
874
+ self.kv_connector.pre_forward(scheduler_output)
875
hidden_states = self.cudagraph_manager.run(
876
input_batch.num_tokens_after_padding
877
)
@@ -877,13 +889,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
889
cudagraph_runtime_mode=CUDAGraphMode.NONE,
890
num_tokens_across_dp=num_tokens_across_dp,
891
):
892
+ self.kv_connector.pre_forward(scheduler_output)
893
hidden_states = self.model(
894
input_ids=input_batch.input_ids,
895
positions=positions,
896
inputs_embeds=input_batch.inputs_embeds,
897
)
898
886
- self.execute_model_state = hidden_states, input_batch
899
+ kv_connector_output = self.kv_connector.post_forward(scheduler_output)
900
+ self.execute_model_state = hidden_states, input_batch, kv_connector_output
901
return None
902
903
@torch.inference_mode()
@@ -892,7 +906,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
906
grammar_output: GrammarOutput | None,
907
) -> AsyncOutput | ModelRunnerOutput:
908
assert self.execute_model_state is not None
895
- hidden_states, input_batch = self.execute_model_state
909
+ hidden_states, input_batch, kv_connector_output = self.execute_model_state
910
self.execute_model_state = None # type: ignore
911
912
sampler_output, num_sampled, num_rejected = self.sample(
@@ -917,6 +931,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
931
req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
932
sampled_token_ids=None, # type: ignore
933
prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
934
+ kv_connector_output=kv_connector_output,
935
)
936
async_output = AsyncOutput(
937
model_runner_output=model_runner_output,
vllm/v1/worker/gpu_worker.py
+4
-2
@@ -24,6 +24,7 @@ from vllm.distributed import (
24
from vllm.distributed.ec_transfer import ensure_ec_transfer_initialized
25
from vllm.distributed.kv_transfer import (
26
ensure_kv_transfer_initialized,
27
+ ensure_kv_transfer_shutdown,
28
get_kv_transfer_group,
29
has_kv_transfer_group,
30
)
@@ -921,8 +922,9 @@ class Worker(WorkerBase):
922
)
923
924
def shutdown(self) -> None:
924
- if runner := getattr(self, "model_runner", None):
925
- runner.ensure_kv_transfer_shutdown()
925
+ # has_kv_transfer_group can be None during interpreter shutdown.
926
+ if ensure_kv_transfer_shutdown is not None:
927
+ ensure_kv_transfer_shutdown()
928
if self.profiler is not None:
929
self.profiler.shutdown()
930