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

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, 23:45 UTC

[KVCache] Make KVCacheSpec hashable (#21791)

Signed-off-by: Chen Zhang <zhangch99@outlook.com>

Chen Zhang committed Jul 29, 2025 at 04:58 UTC 755fa8b657e3666cc93b08a6d2b9a50d0f46c37e

5 files changed +100 -88

tests/v1/core/test_kv_cache_utils.py

+33 -1

@@ -17,7 +17,7 @@ from vllm.v1.core.kv_cache_utils import (

17 estimate_max_model_len, generate_block_hash_extra_keys,

18 get_kv_cache_config, get_max_concurrency_for_kv_cache_config,

19 hash_block_tokens, hash_request_tokens, init_none_hash,

20 - unify_kv_cache_configs)

20 + is_kv_cache_type_uniform, unify_kv_cache_configs)

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

22 KVCacheGroupSpec, KVCacheTensor,

23 SlidingWindowSpec)

@@ -685,6 +685,38 @@ def test_merge_kv_cache_spec():

685 assert merged_layer_spec.sliding_window == 1

686

687

688 +def test_is_kv_cache_type_uniform():

689 + kv_cache_spec = {

690 + "layer_1": new_kv_cache_spec(num_kv_heads=32),

691 + "layer_2": new_kv_cache_spec(num_kv_heads=32),

692 + }

693 + assert is_kv_cache_type_uniform(kv_cache_spec)

694 +

695 + kv_cache_spec = {

696 + "layer_1": new_kv_cache_spec(num_kv_heads=32),

697 + "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1),

698 + }

699 + assert is_kv_cache_type_uniform(kv_cache_spec)

700 +

701 + kv_cache_spec = {

702 + "layer_1": new_kv_cache_spec(num_kv_heads=32),

703 + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),

704 + }

705 + assert not is_kv_cache_type_uniform(kv_cache_spec)

706 +

707 + kv_cache_spec = {

708 + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),

709 + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),

710 + }

711 + assert is_kv_cache_type_uniform(kv_cache_spec)

712 +

713 + kv_cache_spec = {

714 + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1),

715 + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2),

716 + }

717 + assert not is_kv_cache_type_uniform(kv_cache_spec)

718 +

719 +

720 @pytest.mark.parametrize(

721 ("model_id", "max_model_len", "want_estimated_max_len"), [

722 ("Qwen/Qwen1.5-7B", 16385, 16384),

tests/v1/e2e/test_correctness_sliding_window.py

+6 -2

@@ -30,7 +30,9 @@ model_config = {

30 ])

31 @pytest.mark.parametrize("batch_size", [5])

32 @pytest.mark.parametrize("seed", [1])

33 -def test_sliding_window_retrieval(monkeypatch, model, batch_size, seed):

33 +@pytest.mark.parametrize("disable_hybrid_kv_cache_manager", [True, False])

34 +def test_sliding_window_retrieval(monkeypatch, model, batch_size, seed,

35 + disable_hybrid_kv_cache_manager):

36 """

37 The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then

38 asks for value of one of them (which is outside the sliding window).

@@ -42,7 +44,9 @@ def test_sliding_window_retrieval(monkeypatch, model, batch_size, seed):

45 test_config = model_config[model]

45 - llm = LLM(model=model)

47 + llm = LLM(

48 + model=model,

49 + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager)

50 sampling_params = SamplingParams(temperature=0.0, max_tokens=100)

52 prompts, answer, indices = prep_prompts(batch_size,

vllm/v1/core/kv_cache_coordinator.py

+14 -17

@@ -7,7 +7,8 @@ from vllm.v1.core.block_pool import BlockPool

7 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

8 from vllm.v1.core.single_type_kv_cache_manager import (

9 FullAttentionManager, get_manager_for_kv_cache_spec)

10 -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig

10 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

11 + KVCacheSpec)

12 from vllm.v1.request import Request

@@ -258,44 +259,40 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

259 one of them is full attention. Then, split the kv cache groups into full

260 attention groups and other groups.

261 """

261 - full_attention_type_id: Optional[str] = None

262 - other_type_id: Optional[str] = None

262 + full_attention_spec: Optional[FullAttentionSpec] = None

263 + other_spec: Optional[KVCacheSpec] = None

264 self.full_attention_group_ids: list[int] = []

265 self.other_group_ids: list[int] = []

266 for i, g in enumerate(self.kv_cache_config.kv_cache_groups):

267 if isinstance(g.kv_cache_spec, FullAttentionSpec):

267 - if full_attention_type_id is None:

268 - full_attention_type_id = g.kv_cache_spec.type_id

268 + if full_attention_spec is None:

269 + full_attention_spec = g.kv_cache_spec

270 else:

270 - assert full_attention_type_id == g.kv_cache_spec.type_id, (

271 + assert full_attention_spec == g.kv_cache_spec, (

272 "HybridKVCacheCoordinator assumes exactly one type of "

273 "full attention groups now.")

274 self.full_attention_group_ids.append(i)

275 else:

275 - if other_type_id is None:

276 - other_type_id = g.kv_cache_spec.type_id

276 + if other_spec is None:

277 + other_spec = g.kv_cache_spec

278 else:

278 - assert other_type_id == g.kv_cache_spec.type_id, (

279 + assert other_spec == g.kv_cache_spec, (

280 "HybridKVCacheCoordinator assumes "

281 "exactly one other type of groups now.")

282 self.other_group_ids.append(i)

283

283 - assert full_attention_type_id is not None, (

284 + assert full_attention_spec is not None, (

285 "HybridKVCacheCoordinator assumes exactly one type of full "

286 "attention groups now.")

286 - assert other_type_id is not None, (

287 + assert other_spec is not None, (

288 "HybridKVCacheCoordinator assumes exactly one type of other "

289 "groups now.")

290

291 self.full_attention_manager_cls = FullAttentionManager

292 self.other_attention_cls = self.single_type_managers[

293 self.other_group_ids[0]].__class__

293 -

294 - self.full_attention_spec = self.kv_cache_config.kv_cache_groups[

295 - self.full_attention_group_ids[0]].kv_cache_spec

296 - self.other_spec = self.kv_cache_config.kv_cache_groups[

297 - self.other_group_ids[0]].kv_cache_spec

298 -

294 + self.full_attention_spec = full_attention_spec

295 + self.other_spec = other_spec

296 self.full_attention_block_size = self.full_attention_spec.block_size

297 self.other_block_size = self.other_spec.block_size

298

vllm/v1/core/kv_cache_utils.py

+18 -17

@@ -5,7 +5,7 @@

5 import os

6 from collections import defaultdict, deque

7 from collections.abc import Iterable, Sequence

8 -from dataclasses import dataclass

8 +from dataclasses import astuple, dataclass

9 from typing import Any, Callable, NamedTuple, Optional

11 from vllm.config import VllmConfig

@@ -727,7 +727,9 @@ def create_kv_cache_group_specs(

727

728 def is_kv_cache_type_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

729 """

730 - Whether all layers in the given KVCacheSpec have the same type of KV cache.

730 + Whether all layers in the given KVCacheSpec have the same KV cache spec.

731 + Note that we regard FullAttentionSpec with and without sliding window as

732 + the same type.

733

734 Args:

735 kv_cache_spec: The kv cache spec of each attention layer in the model

@@ -736,8 +738,12 @@ def is_kv_cache_type_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

738 True if all layers have the same type, False otherwise.

739 """

740

739 - layer_keys = set(layer.type_id for layer in kv_cache_spec.values())

740 - return len(layer_keys) == 1

741 + try:

742 + kv_cache_spec_values = list(kv_cache_spec.values())

743 + _ = kv_cache_spec_values[0].merge(kv_cache_spec_values)

744 + except AssertionError:

745 + return False

746 + return True

747

748

749 def get_max_concurrency_for_kv_cache_config(

@@ -928,12 +934,12 @@ def _get_kv_cache_config_uniform_page_size(

934 Returns:

935 The generated KVCacheConfig

936 """

931 - # Group all layers by type_id.

937 + # Group all layers by kv_cache_spec.

938 # E.g., 2 full attention layers and 3 sliding window attention layers,

939 # -> (full.0, full.1), (sw.0, sw.1, sw.2).

934 - same_type_layers: dict[str, list[str]] = defaultdict(list)

940 + same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list)

941 for layer_name, layer_spec in kv_cache_spec.items():

936 - same_type_layers[layer_spec.type_id].append(layer_name)

942 + same_type_layers[layer_spec].append(layer_name)

943

944 # Split each group into smaller groups, to make the number of layers in each

945 # group identical. Add padding to the last group of each type if necessary.

@@ -1017,12 +1023,7 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

1023 kv_cache_spec: The kv cache spec of each attention layer in the model

1024 """

1025

1020 - def is_hybrid(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

1021 - type_ids = set(layer_spec.type_id

1022 - for layer_spec in kv_cache_spec.values())

1023 - return len(type_ids) > 1

1024 -

1025 - if not is_hybrid(kv_cache_spec):

1026 + if is_kv_cache_type_uniform(kv_cache_spec):

1027 return

1028

1029 logger.warning(

@@ -1060,7 +1061,7 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

1061 attention_chunk_size=spec.attention_chunk_size,

1062 )

1063

1063 - if is_hybrid(kv_cache_spec):

1064 + if not is_kv_cache_type_uniform(kv_cache_spec):

1065 raise ValueError("Hybrid KV cache manager is disabled but failed to "

1066 "convert the KV cache specs to one unified type.")

1067

@@ -1119,11 +1120,11 @@ def unify_kv_cache_configs(kv_cache_configs: list[KVCacheConfig]):

1120 in-place modified to make them consistent.

1121 """

1122

1122 - # Sort the kv cache groups by the type_id of their KV cache spec.

1123 + # Sort the kv cache groups by their KV cache spec.

1124 # This can avoid the inconsistency caused by the order of groups.

1125 for kv_cache_config in kv_cache_configs:

1125 - kv_cache_config.kv_cache_groups.sort(

1126 - key=lambda x: x.kv_cache_spec.type_id)

1126 + kv_cache_config.kv_cache_groups.sort(key=lambda x: (type(

1127 + x.kv_cache_spec).__name__, astuple(x.kv_cache_spec)))

1128

1129 # Verify that the groups of each rank are the same.

1130 for kv_cache_config in kv_cache_configs[1:]:

vllm/v1/kv_cache_interface.py

+29 -51

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

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

4 import copy

5 -from dataclasses import dataclass

5 +from dataclasses import dataclass, fields

6 from math import prod

7 from typing import Optional

@@ -16,7 +16,7 @@ from vllm.utils import cdiv, get_dtype_size

16 logger = init_logger(__name__)

19 -@dataclass

19 +@dataclass(frozen=True)

20 class KVCacheSpec:

21 """

22 A base class for specifying the KV cache format of one layer.

@@ -25,20 +25,6 @@ class KVCacheSpec:

25 # number of tokens in a block

26 block_size: int

28 - @property

29 - def type_id(self) -> str:

30 - """

31 - The type identifier of this KV cache.

32 - Return different strings for layers with different KV cache type (e.g.,

33 - different number of tokens like full attention vs sliding window

34 - attention, different KV cache size per token like layers with different

35 - number of heads)

36 -

37 - Returns:

38 - The type identifier of this KV cache.

39 - """

40 - raise NotImplementedError

41 -

28 @property

29 def page_size_bytes(self) -> int:

30 """

@@ -63,13 +49,12 @@ class KVCacheSpec:

49 """

50 Merge a list of KVCacheSpec objects into a single KVCacheSpec object.

51 """

66 - assert all(spec.type_id == specs[0].type_id for spec in specs[1:]), (

67 - "All layers in the same KV cache group must share the same "

68 - "type_id.")

52 + assert all(spec == specs[0] for spec in specs[1:]), (

53 + "All layers in the same KV cache group must be the same.")

54 return copy.deepcopy(specs[0])

72 -@dataclass

57 +@dataclass(frozen=True)

58 class AttentionSpec(KVCacheSpec):

59 num_kv_heads: int

60 head_size: int

@@ -84,7 +69,7 @@ class AttentionSpec(KVCacheSpec):

69 * get_dtype_size(self.dtype)

87 -@dataclass

72 +@dataclass(frozen=True)

73 class FullAttentionSpec(AttentionSpec):

74 sliding_window: Optional[int] = None

75 attention_chunk_size: Optional[int] = None

@@ -98,10 +83,6 @@ class FullAttentionSpec(AttentionSpec):

83 Default to None for not using sliding window attention.

84 """

101 - @property

102 - def type_id(self) -> str:

103 - return f"full_attention_{self.block_size}_{self.page_size_bytes}"

104 -

86 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

87 max_model_len = vllm_config.model_config.max_model_len

88 return cdiv(max_model_len, self.block_size) * self.page_size_bytes

@@ -123,15 +104,28 @@ class FullAttentionSpec(AttentionSpec):

104 Merge a list of FullAttentionSpec objects into a single

105 FullAttentionSpec object.

106 """

126 - merged_spec = super().merge(specs)

107 + assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (

108 + "All attention layers in the same KV cache group must be "

109 + "FullAttentionSpec.")

110 +

111 sliding_window = set(spec.sliding_window for spec in specs

112 if spec.sliding_window is not None)

113 attention_chunk_size = set(spec.attention_chunk_size for spec in specs

114 if spec.attention_chunk_size is not None)

131 -

132 - merged_spec.sliding_window = cls.merge_window_sizes(sliding_window)

133 - merged_spec.attention_chunk_size = (

134 - cls.merge_window_sizes(attention_chunk_size))

115 + merged_spec = cls(

116 + block_size=specs[0].block_size,

117 + num_kv_heads=specs[0].num_kv_heads,

118 + head_size=specs[0].head_size,

119 + dtype=specs[0].dtype,

120 + use_mla=specs[0].use_mla,

121 + sliding_window=cls.merge_window_sizes(sliding_window),

122 + attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),

123 + )

124 + for spec in specs:

125 + for f in fields(AttentionSpec):

126 + assert getattr(spec, f.name) == getattr(merged_spec, f.name), (

127 + "All attention layers in the same KV cache group must have "

128 + "the same attention spec.")

129 assert (

130 (merged_spec.sliding_window is not None) +

131 (merged_spec.attention_chunk_size is not None) <= 1

@@ -140,16 +134,10 @@ class FullAttentionSpec(AttentionSpec):

134 return merged_spec

135

136

143 -@dataclass

137 +@dataclass(frozen=True)

138 class ChunkedLocalAttentionSpec(AttentionSpec):

139 attention_chunk_size: int

140

147 - @property

148 - def type_id(self) -> str:

149 - return (

150 - f"local_attention_{self.attention_chunk_size}_{self.block_size}_{self.page_size_bytes}"

151 - ) # noqa

152 -

141 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

142 max_model_len = vllm_config.model_config.max_model_len

143 max_num_batched_tokens = (

@@ -165,17 +153,13 @@ class ChunkedLocalAttentionSpec(AttentionSpec):

153 return cdiv(num_tokens, self.block_size) * self.page_size_bytes

154

155

168 -@dataclass

156 +@dataclass(frozen=True)

157 class SlidingWindowSpec(AttentionSpec):

158 sliding_window: int

159

160 def __post_init__(self):

161 assert not self.use_mla, "MLA is not supported for sliding window"

162

175 - @property

176 - def type_id(self) -> str:

177 - return f"sliding_window_{self.sliding_window}_{self.block_size}_{self.page_size_bytes}" # noqa

178 -

163 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

164 max_model_len = vllm_config.model_config.max_model_len

165 max_num_batched_tokens = (

@@ -195,23 +179,17 @@ class SlidingWindowSpec(AttentionSpec):

179 return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes

180

181

198 -@dataclass

182 +@dataclass(frozen=True)

183 class MambaSpec(KVCacheSpec):

184 shapes: tuple[tuple[int, ...], ...]

185 dtype: torch.dtype

186 page_size_padded: Optional[int] = None

187 mamba_type: str = "mamba2"

188

205 - def __post_init__(self):

206 - self.num_elements = sum(prod(shape) for shape in self.shapes)

207 -

208 - @property

209 - def type_id(self) -> str:

210 - return f"mamba_{self.shapes}_{self.dtype}_{self.mamba_type}"

211 -

189 @property

190 def page_size_bytes(self) -> int:

214 - page_size = self.num_elements * get_dtype_size(self.dtype)

191 + num_elements = sum(prod(shape) for shape in self.shapes)

192 + page_size = num_elements * get_dtype_size(self.dtype)

193 if self.page_size_padded is not None:

194 assert self.page_size_padded >= page_size

195 return self.page_size_padded