From 47d4407d7cbb2a70fd02d73f34407caeee6ae565 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 22 May 2026 15:18:23 -0700 Subject: [PATCH] [Model Runner V2] Support sharing kv cache layers (#35045) Signed-off-by: Nick Hill --- vllm/config/vllm.py | 13 ++++---- vllm/v1/worker/gpu/attn_utils.py | 52 ++++++++++++++++++++++++------ vllm/v1/worker/gpu/model_runner.py | 5 ++- vllm/v1/worker/utils.py | 3 ++ 4 files changed, 54 insertions(+), 19 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 20309f0c9f..2c020e9cff 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -8,6 +8,7 @@ import os import tempfile import threading import time +from collections.abc import Iterable from contextlib import contextmanager from dataclasses import is_dataclass from datetime import datetime @@ -2220,7 +2221,7 @@ T = TypeVar("T") def get_layers_from_vllm_config( vllm_config: VllmConfig, layer_type: type[T], - layer_names: list[str] | None = None, + layer_names: Iterable[str] | None = None, ) -> dict[str, T]: """ Get layers from the vLLM config. @@ -2231,14 +2232,12 @@ def get_layers_from_vllm_config( layer_names: The names of the layers to get. If None, return all layers. """ - if layer_names is None: - layer_names = list(vllm_config.compilation_config.static_forward_context.keys()) - forward_context = vllm_config.compilation_config.static_forward_context + if layer_names is None: + layer_names = forward_context.keys() return { - layer_name: forward_context[layer_name] + layer_name: layer for layer_name in layer_names - if layer_name in forward_context - and isinstance(forward_context[layer_name], layer_type) + if isinstance(layer := forward_context.get(layer_name), layer_type) } diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 53ff79eb49..8b4e0c3c0a 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -7,6 +7,7 @@ from typing import Any, cast import torch from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import ( @@ -24,6 +25,7 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.utils import ( AttentionGroup, + add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache, prepare_kernel_block_sizes, ) @@ -40,12 +42,24 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: layer_type = cast(type[Any], AttentionLayerBase) attn_layers = get_layers_from_vllm_config(vllm_config, layer_type) for layer_name, attn_module in attn_layers.items(): + if getattr(attn_module, "kv_sharing_target_layer_name", None): + # This layer will use KV cache of the sharing target layer. + continue # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): kv_cache_spec[layer_name] = spec return kv_cache_spec +def get_shared_kv_cache_layers(vllm_config: VllmConfig): + attn_layers = get_layers_from_vllm_config(vllm_config, Attention) + return { + layer_name: kv_tgt_layer + for layer_name, attn_module in attn_layers.items() + if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name) + } + + def init_attn_backend( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, @@ -60,6 +74,12 @@ def init_attn_backend( attn_backends: dict[str, type[AttentionBackend]] = {} attn_groups: list[list[AttentionGroup]] = [] + # Add KV-sharing layers to their target's kv cache group so they are + # discovered alongside the target layer in Phase 1 below. + add_kv_sharing_layers_to_kv_cache_groups( + get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups + ) + # Phase 1: discover attention groups for each kv cache group. for kv_cache_group_id, kv_cache_group_spec in enumerate( kv_cache_config.kv_cache_groups @@ -85,10 +105,7 @@ def init_attn_backend( key = (attn_backend.full_cls_name(), layer_kv_cache_spec) if key not in group_map: group_map[key] = AttentionGroup( - attn_backend, - [layer_name], - layer_kv_cache_spec, - kv_cache_group_id, + attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id ) group_order.append(key) else: @@ -143,7 +160,9 @@ def init_attn_backend( ) -def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device): +def _allocate_kv_cache( + kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device +): kv_cache_raw_tensors: dict[str, torch.Tensor] = {} for kv_cache_tensor in kv_cache_config.kv_cache_tensors: 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): for group in kv_cache_config.kv_cache_groups: for layer_name in group.layer_names: layer_names.add(layer_name) - assert layer_names == set(kv_cache_raw_tensors.keys()), ( + assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), ( "Some layers are not correctly initialized" ) return kv_cache_raw_tensors @@ -166,6 +185,7 @@ def _reshape_kv_cache( attn_backends: dict[str, type[AttentionBackend]], cache_dtype: str, kernel_block_sizes: list[int], + shared_kv_cache_layers: dict[str, str], ) -> dict[str, Any]: kv_caches: dict[str, Any] = {} has_attn, has_mamba = False, False @@ -173,6 +193,9 @@ def _reshape_kv_cache( kv_cache_config.kv_cache_groups ): for layer_name in kv_cache_group_spec.layer_names: + if layer_name in shared_kv_cache_layers: + # Shared layer — tensor will be aliased to its target later. + continue kv_cache_spec = kv_cache_group_spec.kv_cache_spec if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] @@ -270,15 +293,21 @@ def _reshape_kv_cache( if has_attn and has_mamba: _update_hybrid_attention_layout(kv_caches, kv_cache_config) + # Map any sharing layers to their target layer's KV cache. + for layer_name, target_layer_name in shared_kv_cache_layers.items(): + kv_caches[layer_name] = kv_caches[target_layer_name] + return kv_caches def _update_hybrid_attention_layout( - kv_caches: dict[str, Any], - kv_cache_config: KVCacheConfig, + kv_caches: dict[str, Any], kv_cache_config: KVCacheConfig ) -> None: for kv_cache_group_spec in kv_cache_config.kv_cache_groups: for layer_name in kv_cache_group_spec.layer_names: + if layer_name not in kv_caches: + # Shared layer — will be aliased to its target after this pass. + continue kv_cache_spec = kv_cache_group_spec.kv_cache_spec if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] @@ -308,14 +337,19 @@ def init_kv_cache( device: torch.device, cache_dtype: str, kernel_block_sizes: list[int], + vllm_config: VllmConfig, ) -> dict[str, Any]: - kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device) + shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config) + kv_cache_raw_tensors = _allocate_kv_cache( + kv_cache_config, shared_kv_cache_layers, device + ) kv_caches = _reshape_kv_cache( kv_cache_config, kv_cache_raw_tensors, attn_backends, cache_dtype, kernel_block_sizes, + shared_kv_cache_layers, ) bind_kv_cache(kv_caches, forward_context, runner_kv_caches) return kv_caches diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3b965dfc6e..c0a95618c1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -432,9 +432,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.speculator is not None: # HACK(woosuk) self.speculator.set_attn( - self.model_state, - self.kv_cache_config, - self.block_tables, + self.model_state, self.kv_cache_config, self.block_tables ) self.kv_caches: list[torch.Tensor] = [] @@ -446,6 +444,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.device, self.cache_config.cache_dtype, kernel_block_sizes, + self.vllm_config, ) self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b693adbf27..7cb1620c95 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -441,6 +441,9 @@ def add_kv_sharing_layers_to_kv_cache_groups( from the KV cache of `shared_kv_cache_layers[layer_name]`. kv_cache_groups: The KV cache groups of the model. """ + if not shared_kv_cache_layers: + return + layer_to_kv_cache_group: dict[str, KVCacheGroupSpec] = {} for kv_cache_group in kv_cache_groups: for layer_name in kv_cache_group.layer_names: