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 3, 00:30 UTC
[Model Runner V2] Support sharing kv cache layers (#35045)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 22, 2026 at 15:18 UTC
47d4407d7cbb2a70fd02d73f34407caeee6ae565
4 files changed
+54
-19
vllm/config/vllm.py
+6
-7
@@ -8,6 +8,7 @@ import os
8
import tempfile
9
import threading
10
import time
11
+from collections.abc import Iterable
12
from contextlib import contextmanager
13
from dataclasses import is_dataclass
14
from datetime import datetime
@@ -2220,7 +2221,7 @@ T = TypeVar("T")
2221
def get_layers_from_vllm_config(
2222
vllm_config: VllmConfig,
2223
layer_type: type[T],
2223
- layer_names: list[str] | None = None,
2224
+ layer_names: Iterable[str] | None = None,
2225
) -> dict[str, T]:
2226
"""
2227
Get layers from the vLLM config.
@@ -2231,14 +2232,12 @@ def get_layers_from_vllm_config(
2232
layer_names: The names of the layers to get. If None, return all layers.
2233
"""
2234
2234
- if layer_names is None:
2235
- layer_names = list(vllm_config.compilation_config.static_forward_context.keys())
2236
-
2235
forward_context = vllm_config.compilation_config.static_forward_context
2236
+ if layer_names is None:
2237
+ layer_names = forward_context.keys()
2238
2239
return {
2240
- layer_name: forward_context[layer_name]
2240
+ layer_name: layer
2241
for layer_name in layer_names
2242
- if layer_name in forward_context
2243
- and isinstance(forward_context[layer_name], layer_type)
2242
+ if isinstance(layer := forward_context.get(layer_name), layer_type)
2243
}
vllm/v1/worker/gpu/attn_utils.py
+43
-9
@@ -7,6 +7,7 @@ from typing import Any, cast
7
import torch
8
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
10
+from vllm.model_executor.layers.attention import Attention
11
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12
from vllm.utils.torch_utils import get_dtype_size
13
from vllm.v1.attention.backend import (
@@ -24,6 +25,7 @@ from vllm.v1.kv_cache_interface import (
25
from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
26
from vllm.v1.worker.utils import (
27
AttentionGroup,
28
+ add_kv_sharing_layers_to_kv_cache_groups,
29
bind_kv_cache,
30
prepare_kernel_block_sizes,
31
)
@@ -40,12 +42,24 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
42
layer_type = cast(type[Any], AttentionLayerBase)
43
attn_layers = get_layers_from_vllm_config(vllm_config, layer_type)
44
for layer_name, attn_module in attn_layers.items():
45
+ if getattr(attn_module, "kv_sharing_target_layer_name", None):
46
+ # This layer will use KV cache of the sharing target layer.
47
+ continue
48
# Skip modules that don't need KV cache (eg encoder-only attention)
49
if spec := attn_module.get_kv_cache_spec(vllm_config):
50
kv_cache_spec[layer_name] = spec
51
return kv_cache_spec
52
53
54
+def get_shared_kv_cache_layers(vllm_config: VllmConfig):
55
+ attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
56
+ return {
57
+ layer_name: kv_tgt_layer
58
+ for layer_name, attn_module in attn_layers.items()
59
+ if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name)
60
+ }
61
+
62
+
63
def init_attn_backend(
64
kv_cache_config: KVCacheConfig,
65
vllm_config: VllmConfig,
@@ -60,6 +74,12 @@ def init_attn_backend(
74
attn_backends: dict[str, type[AttentionBackend]] = {}
75
attn_groups: list[list[AttentionGroup]] = []
76
77
+ # Add KV-sharing layers to their target's kv cache group so they are
78
+ # discovered alongside the target layer in Phase 1 below.
79
+ add_kv_sharing_layers_to_kv_cache_groups(
80
+ get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups
81
+ )
82
+
83
# Phase 1: discover attention groups for each kv cache group.
84
for kv_cache_group_id, kv_cache_group_spec in enumerate(
85
kv_cache_config.kv_cache_groups
@@ -85,10 +105,7 @@ def init_attn_backend(
105
key = (attn_backend.full_cls_name(), layer_kv_cache_spec)
106
if key not in group_map:
107
group_map[key] = AttentionGroup(
88
- attn_backend,
89
- [layer_name],
90
- layer_kv_cache_spec,
91
- kv_cache_group_id,
108
+ attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id
109
)
110
group_order.append(key)
111
else:
@@ -143,7 +160,9 @@ def init_attn_backend(
160
)
161
162
146
-def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
163
+def _allocate_kv_cache(
164
+ kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device
165
+):
166
kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
167
for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
168
tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
@@ -154,7 +173,7 @@ def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
173
for group in kv_cache_config.kv_cache_groups:
174
for layer_name in group.layer_names:
175
layer_names.add(layer_name)
157
- assert layer_names == set(kv_cache_raw_tensors.keys()), (
176
+ assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), (
177
"Some layers are not correctly initialized"
178
)
179
return kv_cache_raw_tensors
@@ -166,6 +185,7 @@ def _reshape_kv_cache(
185
attn_backends: dict[str, type[AttentionBackend]],
186
cache_dtype: str,
187
kernel_block_sizes: list[int],
188
+ shared_kv_cache_layers: dict[str, str],
189
) -> dict[str, Any]:
190
kv_caches: dict[str, Any] = {}
191
has_attn, has_mamba = False, False
@@ -173,6 +193,9 @@ def _reshape_kv_cache(
193
kv_cache_config.kv_cache_groups
194
):
195
for layer_name in kv_cache_group_spec.layer_names:
196
+ if layer_name in shared_kv_cache_layers:
197
+ # Shared layer — tensor will be aliased to its target later.
198
+ continue
199
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
200
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
201
kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
@@ -270,15 +293,21 @@ def _reshape_kv_cache(
293
if has_attn and has_mamba:
294
_update_hybrid_attention_layout(kv_caches, kv_cache_config)
295
296
+ # Map any sharing layers to their target layer's KV cache.
297
+ for layer_name, target_layer_name in shared_kv_cache_layers.items():
298
+ kv_caches[layer_name] = kv_caches[target_layer_name]
299
+
300
return kv_caches
301
302
303
def _update_hybrid_attention_layout(
277
- kv_caches: dict[str, Any],
278
- kv_cache_config: KVCacheConfig,
304
+ kv_caches: dict[str, Any], kv_cache_config: KVCacheConfig
305
) -> None:
306
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
307
for layer_name in kv_cache_group_spec.layer_names:
308
+ if layer_name not in kv_caches:
309
+ # Shared layer — will be aliased to its target after this pass.
310
+ continue
311
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
312
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
313
kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
@@ -308,14 +337,19 @@ def init_kv_cache(
337
device: torch.device,
338
cache_dtype: str,
339
kernel_block_sizes: list[int],
340
+ vllm_config: VllmConfig,
341
) -> dict[str, Any]:
312
- kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
342
+ shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
343
+ kv_cache_raw_tensors = _allocate_kv_cache(
344
+ kv_cache_config, shared_kv_cache_layers, device
345
+ )
346
kv_caches = _reshape_kv_cache(
347
kv_cache_config,
348
kv_cache_raw_tensors,
349
attn_backends,
350
cache_dtype,
351
kernel_block_sizes,
352
+ shared_kv_cache_layers,
353
)
354
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
355
return kv_caches
vllm/v1/worker/gpu/model_runner.py
+2
-3
@@ -432,9 +432,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
432
if self.speculator is not None:
433
# HACK(woosuk)
434
self.speculator.set_attn(
435
- self.model_state,
436
- self.kv_cache_config,
437
- self.block_tables,
435
+ self.model_state, self.kv_cache_config, self.block_tables
436
)
437
438
self.kv_caches: list[torch.Tensor] = []
@@ -446,6 +444,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
444
self.device,
445
self.cache_config.cache_dtype,
446
kernel_block_sizes,
447
+ self.vllm_config,
448
)
449
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
450
vllm/v1/worker/utils.py
+3
@@ -441,6 +441,9 @@ def add_kv_sharing_layers_to_kv_cache_groups(
441
from the KV cache of `shared_kv_cache_layers[layer_name]`.
442
kv_cache_groups: The KV cache groups of the model.
443
"""
444
+ if not shared_kv_cache_layers:
445
+ return
446
+
447
layer_to_kv_cache_group: dict[str, KVCacheGroupSpec] = {}
448
for kv_cache_group in kv_cache_groups:
449
for layer_name in kv_cache_group.layer_names: