iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model Runner V2] Support sharing kv cache layers (#35045)

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Nick Hill
2026-05-22 22:18:23 +00:00
committed by GitHub
parent e203006a8b
commit 47d4407d7c
4 changed files with 54 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/config/vllm.py
+6 -7
View File
@@ -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)
}
vllm/v1/worker/gpu/attn_utils.py
+43 -9
View File
@@ -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
vllm/v1/worker/gpu/model_runner.py
+2 -3
View File
@@ -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)
vllm/v1/worker/utils.py
+3
View File
@@ -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: