@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9607d5eb44

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, 05:30 UTC

[Hybrid Allocator] Support full attention with different hidden size (#25101)

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

Chen Zhang committed Sep 19, 2025 at 23:43 UTC 9607d5eb449711b349d4c2bee0a9c94afcc7ed14

6 files changed +325 -93

tests/v1/core/test_kv_cache_utils.py

+103 -15

@@ -18,12 +18,14 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager

18 from vllm.v1.core.kv_cache_utils import (

19 BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics,

20 estimate_max_model_len, generate_block_hash_extra_keys,

21 - get_kv_cache_configs, get_max_concurrency_for_kv_cache_config,

22 - get_request_block_hasher, hash_block_tokens, init_none_hash,

23 - is_kv_cache_type_uniform, make_block_hash_with_group_id)

21 + generate_scheduler_kv_cache_config, get_kv_cache_configs,

22 + get_max_concurrency_for_kv_cache_config, get_request_block_hasher,

23 + hash_block_tokens, init_none_hash, is_kv_cache_spec_uniform,

24 + make_block_hash_with_group_id)

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

26 KVCacheGroupSpec, KVCacheSpec,

26 - KVCacheTensor, SlidingWindowSpec)

27 + KVCacheTensor, SlidingWindowSpec,

28 + UniformTypeKVCacheSpecs)

29 from vllm.v1.metrics.stats import PrefixCacheStats

30 from vllm.v1.request import Request

@@ -927,36 +929,36 @@ def test_merge_kv_cache_spec():

929 assert merged_layer_spec.sliding_window == 1

930

931

930 -def test_is_kv_cache_type_uniform():

932 +def test_is_kv_cache_spec_uniform():

933 kv_cache_spec = {

934 "layer_1": new_kv_cache_spec(num_kv_heads=32),

935 "layer_2": new_kv_cache_spec(num_kv_heads=32),

936 }

935 - assert is_kv_cache_type_uniform(kv_cache_spec)

937 + assert is_kv_cache_spec_uniform(kv_cache_spec)

938

939 kv_cache_spec = {

940 "layer_1": new_kv_cache_spec(num_kv_heads=32),

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

942 }

941 - assert is_kv_cache_type_uniform(kv_cache_spec)

943 + assert is_kv_cache_spec_uniform(kv_cache_spec)

944

945 kv_cache_spec = {

946 "layer_1": new_kv_cache_spec(num_kv_heads=32),

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

948 }

947 - assert not is_kv_cache_type_uniform(kv_cache_spec)

949 + assert not is_kv_cache_spec_uniform(kv_cache_spec)

950

951 kv_cache_spec = {

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

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

954 }

953 - assert is_kv_cache_type_uniform(kv_cache_spec)

955 + assert is_kv_cache_spec_uniform(kv_cache_spec)

956

957 kv_cache_spec = {

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

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

960 }

959 - assert not is_kv_cache_type_uniform(kv_cache_spec)

961 + assert not is_kv_cache_spec_uniform(kv_cache_spec)

962

963

964 @pytest.mark.parametrize(

@@ -1286,14 +1288,28 @@ def test_get_kv_cache_config_one_worker():

1288 ],

1289 )

1290

1289 - # different hidden size, unimplemented

1291 + # different hidden size

1292 kv_cache_specs_hybrid = {

1293 'layer_1': new_kv_cache_spec(head_size=128),

1292 - 'layer_2': new_kv_cache_spec(),

1294 + 'layer_2': new_kv_cache_spec(head_size=64),

1295 }

1294 - with pytest.raises(NotImplementedError):

1295 - get_kv_cache_configs(vllm_config, [kv_cache_specs_hybrid],

1296 - [mem_per_block_per_layer * 2 * 32])[0]

1296 + kv_cache_config_hybrid = get_kv_cache_configs(

1297 + vllm_config, [kv_cache_specs_hybrid],

1298 + [mem_per_block_per_layer * 3 * 32])[0]

1299 + assert kv_cache_config_hybrid == KVCacheConfig(

1300 + num_blocks=32,

1301 + kv_cache_tensors=[

1302 + KVCacheTensor(size=mem_per_block_per_layer * 32 * 2,

1303 + shared_by=["layer_1"]),

1304 + KVCacheTensor(size=mem_per_block_per_layer * 32,

1305 + shared_by=["layer_2"]),

1306 + ],

1307 + kv_cache_groups=[

1308 + KVCacheGroupSpec(["layer_1", "layer_2"],

1309 + UniformTypeKVCacheSpecs(

1310 + block_size=16,

1311 + kv_cache_specs=kv_cache_specs_hybrid))

1312 + ])

1313

1314 # Test num_gpu_blocks_override

1315 vllm_config.cache_config.num_gpu_blocks_override = 16

@@ -1324,3 +1340,75 @@ def test_get_kv_cache_configs_attention_free():

1340 kv_cache_groups=[],

1341 )

1342 ]

1343 +

1344 +

1345 +def test_generate_uniform_type_kv_cache_specs():

1346 + # All layers are full attention, can be merged

1347 + kv_cache_specs = {

1348 + 'layer_1': new_kv_cache_spec(),

1349 + 'layer_2': new_kv_cache_spec(head_size=128),

1350 + }

1351 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)

1352 + assert uniform_spec == UniformTypeKVCacheSpecs(

1353 + block_size=16, kv_cache_specs=kv_cache_specs)

1354 +

1355 + # Full attention + sliding window, cannot be merged

1356 + kv_cache_specs = {

1357 + 'layer_1': new_kv_cache_spec(),

1358 + 'layer_2': new_sliding_window_spec(sliding_window=1),

1359 + }

1360 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)

1361 + assert uniform_spec is None

1362 +

1363 + # different order of full attention + sliding window, cannot be merged

1364 + kv_cache_specs = {

1365 + 'layer_1': new_sliding_window_spec(sliding_window=1),

1366 + 'layer_2': new_kv_cache_spec(),

1367 + }

1368 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)

1369 + assert uniform_spec is None

1370 +

1371 + # Same-size sliding window, can be merged

1372 + kv_cache_specs = {

1373 + 'layer_1': new_sliding_window_spec(sliding_window=1),

1374 + 'layer_2': new_sliding_window_spec(sliding_window=1, head_size=128),

1375 + }

1376 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)

1377 + assert uniform_spec == UniformTypeKVCacheSpecs(

1378 + block_size=16, kv_cache_specs=kv_cache_specs)

1379 +

1380 + # different block sizes, cannot be merged

1381 + kv_cache_specs = {

1382 + 'layer_1': new_kv_cache_spec(block_size=16),

1383 + 'layer_2': new_kv_cache_spec(block_size=32),

1384 + }

1385 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(kv_cache_specs)

1386 + assert uniform_spec is None

1387 +

1388 +

1389 +def test_generate_scheduler_kv_cache_config():

1390 + kv_cache_specs = {

1391 + 'layer_1': new_kv_cache_spec(),

1392 + 'layer_2': new_kv_cache_spec(head_size=128),

1393 + }

1394 + kv_cache_configs = [

1395 + KVCacheConfig(

1396 + num_blocks=10,

1397 + kv_cache_tensors=[],

1398 + kv_cache_groups=[

1399 + KVCacheGroupSpec(['layer_1', 'layer_2'],

1400 + UniformTypeKVCacheSpecs(

1401 + block_size=16,

1402 + kv_cache_specs=kv_cache_specs)),

1403 + ],

1404 + )

1405 + ]

1406 + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(

1407 + kv_cache_configs)

1408 + assert scheduler_kv_cache_config == KVCacheConfig(

1409 + num_blocks=10,

1410 + kv_cache_tensors=[],

1411 + kv_cache_groups=[

1412 + KVCacheGroupSpec(['layer_1', 'layer_2'], new_kv_cache_spec())

1413 + ],

1414 + )

vllm/v1/core/kv_cache_utils.py

+108 -38

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

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

3 """KV-Cache Utilities."""

5 +import copy

6 import os

7 from collections import defaultdict, deque

8 from collections.abc import Iterable, Sequence

@@ -15,7 +16,8 @@ from vllm.utils import GiB_bytes, cdiv, sha256_cbor

16 from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,

17 FullAttentionSpec, KVCacheConfig,

18 KVCacheGroupSpec, KVCacheSpec,

18 - KVCacheTensor, SlidingWindowSpec)

19 + KVCacheTensor, SlidingWindowSpec,

20 + UniformTypeKVCacheSpecs)

21 from vllm.v1.metrics.stats import PrefixCacheStats

22 from vllm.v1.request import Request

@@ -750,7 +752,7 @@ def create_kv_cache_group_specs(

752 return kv_cache_groups

753

754

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

755 +def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

756 """

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

758 Note that we regard FullAttentionSpec with and without sliding window as

@@ -793,6 +795,21 @@ def get_max_concurrency_for_kv_cache_config(

795 return max_concurrency

796

797

798 +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:

799 + """

800 + Override the number of kv cache blocks if `num_gpu_blocks_override` is set.

801 + """

802 + if vllm_config.cache_config.num_gpu_blocks_override is not None:

803 + num_gpu_blocks_override = \

804 + vllm_config.cache_config.num_gpu_blocks_override

805 + logger.info(

806 + "Overriding num_gpu_blocks=%d with "

807 + "num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)

808 + num_blocks = num_gpu_blocks_override

809 +

810 + return num_blocks

811 +

812 +

813 def get_num_blocks(vllm_config: VllmConfig, num_layers: int,

814 available_memory: int, page_size: int) -> int:

815 """

@@ -806,13 +823,7 @@ def get_num_blocks(vllm_config: VllmConfig, num_layers: int,

823 """

824 num_blocks = int(available_memory // page_size // num_layers)

825 num_blocks = max(num_blocks, 0)

809 - if vllm_config.cache_config.num_gpu_blocks_override is not None:

810 - num_gpu_blocks_override = \

811 - vllm_config.cache_config.num_gpu_blocks_override

812 - logger.info(

813 - "Overriding num_gpu_blocks=%d with "

814 - "num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)

815 - num_blocks = num_gpu_blocks_override

826 + num_blocks = may_override_num_blocks(vllm_config, num_blocks)

827 return num_blocks

828

829

@@ -825,11 +836,11 @@ def get_uniform_page_size(kv_cache_spec: dict[str, KVCacheSpec]) -> int:

836 return page_sizes.pop()

837

838

828 -def _get_kv_cache_groups_uniform_type(

839 +def _get_kv_cache_groups_uniform_spec(

840 kv_cache_specs: dict[str, KVCacheSpec]) -> list[KVCacheGroupSpec]:

841 """

831 - Generates the KV cache configuration for a model with one type of KV cache.

832 - Divide the available memory equally among all layers.

842 + Generates the KV cache configuration for a model with the same KV cache

843 + spec for all layers.

844

845 Args:

846 kv_cache_specs: The kv cache spec of each attention layer in the model

@@ -842,6 +853,22 @@ def _get_kv_cache_groups_uniform_type(

853 [list(kv_cache_specs.keys())])

854

855

856 +def _get_kv_cache_groups_uniform_type(

857 + spec: UniformTypeKVCacheSpecs) -> list[KVCacheGroupSpec]:

858 + """

859 + Generates the KV cache configuration for a model with one type of KV cache

860 + but different hidden sizes. All layers are merged into one group.

861 +

862 + Args:

863 + spec: The UniformTypeKVCacheSpecs of the model

864 +

865 + Returns:

866 + The generated KVCacheGroupSpecs

867 + """

868 +

869 + return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]

870 +

871 +

872 def is_kv_cache_page_size_uniform(

873 kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

874 """

@@ -1000,28 +1027,45 @@ def get_kv_cache_config_from_groups(vllm_config: VllmConfig,

1027 )

1028

1029 # Determine how model runners should initialize the KV cache tensors.

1003 - # We will have group_size memory pools, each is shared by one layer from

1004 - # each group. As layers of different groups have different block table,

1005 - # they will use different parts of the shared Tensor.

1006 - # The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),

1007 - # (sw.1, padding) will be: (group_size = 2)

1008 - # full.0, sw.0, sw.1: share a Tensor with size=available_memory//2

1009 - # full.1, sw.2: share another Tensor with size=available_memory//2

1010 - group_size = max(len(group.layer_names) for group in kv_cache_groups)

1011 -

1012 - page_size = get_uniform_page_size(kv_cache_specs)

1013 - assert group_size > 0, "group_size must be greater than 0"

1014 - num_blocks = get_num_blocks(vllm_config, group_size, available_memory,

1015 - page_size)

1016 - per_memory_pool_size = page_size * num_blocks

1017 - kv_cache_tensors = []

1018 - for i in range(group_size):

1019 - shared_by = []

1020 - for j in range(len(kv_cache_groups)):

1021 - if i < len(kv_cache_groups[j].layer_names):

1022 - shared_by.append(kv_cache_groups[j].layer_names[i])

1023 - kv_cache_tensors.append(

1024 - KVCacheTensor(size=per_memory_pool_size, shared_by=shared_by))

1030 + if len(kv_cache_groups) == 1 and \

1031 + isinstance(kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs):

1032 + # Special case: all layers have the same type of KV cache but with

1033 + # different hidden size. Allocate different amount of memory for each

1034 + # layer based on its hidden size.

1035 + num_blocks = available_memory // kv_cache_groups[

1036 + 0].kv_cache_spec.page_size_bytes

1037 + num_blocks = may_override_num_blocks(vllm_config, num_blocks)

1038 + per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs

1039 + kv_cache_tensors = [

1040 + KVCacheTensor(size=per_layer_specs[layer_name].page_size_bytes *

1041 + num_blocks,

1042 + shared_by=[layer_name])

1043 + for layer_name in kv_cache_groups[0].layer_names

1044 + ]

1045 + else:

1046 + # General case:

1047 + # We will have group_size memory pools, each is shared by one layer from

1048 + # each group. As layers of different groups have different block table,

1049 + # they will use different parts of the shared Tensor.

1050 + # The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),

1051 + # (sw.1, padding) will be: (group_size = 2)

1052 + # full.0, sw.0, sw.1: share a Tensor with size=available_memory//2

1053 + # full.1, sw.2: share another Tensor with size=available_memory//2

1054 + group_size = max(len(group.layer_names) for group in kv_cache_groups)

1055 +

1056 + page_size = get_uniform_page_size(kv_cache_specs)

1057 + assert group_size > 0, "group_size must be greater than 0"

1058 + num_blocks = get_num_blocks(vllm_config, group_size, available_memory,

1059 + page_size)

1060 + kv_cache_tensors = []

1061 + for i in range(group_size):

1062 + shared_by = []

1063 + for j in range(len(kv_cache_groups)):

1064 + if i < len(kv_cache_groups[j].layer_names):

1065 + shared_by.append(kv_cache_groups[j].layer_names[i])

1066 + kv_cache_tensors.append(

1067 + KVCacheTensor(size=page_size * num_blocks,

1068 + shared_by=shared_by))

1069

1070 kv_cache_config = KVCacheConfig(

1071 num_blocks=num_blocks,

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

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

1104 """

1105

1062 - if is_kv_cache_type_uniform(kv_cache_spec):

1106 + if is_kv_cache_spec_uniform(kv_cache_spec):

1107 return

1108

1109 logger.warning(

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

1141 attention_chunk_size=spec.attention_chunk_size,

1142 )

1143

1100 - if not is_kv_cache_type_uniform(kv_cache_spec):

1144 + if not is_kv_cache_spec_uniform(kv_cache_spec):

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

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

1147

@@ -1122,11 +1166,16 @@ def get_kv_cache_groups(

1166 # This returns an empty list to allow for the KVCacheManager to handle

1167 # attention free models.

1168 return []

1125 - elif is_kv_cache_type_uniform(kv_cache_spec):

1169 + elif is_kv_cache_spec_uniform(kv_cache_spec):

1170 # KV cache of all layers are the same, which is true for

1171 # most models. Allocate the same amount of memory for

1172 # each layer.

1129 - return _get_kv_cache_groups_uniform_type(kv_cache_spec)

1173 + return _get_kv_cache_groups_uniform_spec(kv_cache_spec)

1174 + elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec):

1175 + # All layers need the same number of token slots (e.g., all layers are

1176 + # full attention, or all layers are sliding window attention with the

1177 + # same window size). Put all layers into one group.

1178 + return _get_kv_cache_groups_uniform_type(uniform_spec)

1179 elif is_kv_cache_page_size_uniform(kv_cache_spec):

1180 # Model contains multiple attention types, but KV cache of all layers

1181 # have the same physical memory per block per layer. Split the layers

@@ -1137,6 +1186,27 @@ def get_kv_cache_groups(

1186 raise NotImplementedError

1187

1188

1189 +def generate_scheduler_kv_cache_config(

1190 + kv_cache_configs: list[KVCacheConfig]) -> KVCacheConfig:

1191 + """

1192 + Generate the KV cache configuration for the scheduler.

1193 + """

1194 + assert all([

1195 + cfg.num_blocks == kv_cache_configs[0].num_blocks

1196 + for cfg in kv_cache_configs

1197 + ])

1198 + # All workers have the same kv_cache_config except layer names, so use

1199 + # an arbitrary one to initialize the scheduler.

1200 + cfg = copy.deepcopy(kv_cache_configs[0])

1201 + for group in cfg.kv_cache_groups:

1202 + if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):

1203 + # All layers in the UniformTypeKVCacheSpecs have the same type,

1204 + # so use an arbitrary one to initialize the scheduler.

1205 + group.kv_cache_spec = next(

1206 + iter(group.kv_cache_spec.kv_cache_specs.values()))

1207 + return cfg

1208 +

1209 +

1210 def get_kv_cache_configs(vllm_config: VllmConfig,

1211 kv_cache_specs: list[dict[str, KVCacheSpec]],

1212 available_memory: list[int]) -> list[KVCacheConfig]:

vllm/v1/engine/core.py

+6 -10

@@ -29,7 +29,9 @@ from vllm.transformers_utils.config import (

29 maybe_register_config_serialize_by_value)

30 from vllm.utils import (decorate_logs, get_hash_fn_by_name, make_zmq_socket,

31 resolve_obj_by_qualname, set_process_title)

32 -from vllm.v1.core.kv_cache_utils import (BlockHash, get_kv_cache_configs,

32 +from vllm.v1.core.kv_cache_utils import (BlockHash,

33 + generate_scheduler_kv_cache_config,

34 + get_kv_cache_configs,

35 get_request_block_hasher,

36 init_none_hash)

37 from vllm.v1.core.sched.interface import SchedulerInterface

@@ -196,16 +198,10 @@ class EngineCore:

198

199 kv_cache_configs = get_kv_cache_configs(vllm_config, kv_cache_specs,

200 available_gpu_memory)

199 -

200 - # All workers have the same kv_cache_config except layer names, so use

201 - # an arbitrary one to initialize the scheduler.

202 - assert all([

203 - cfg.num_blocks == kv_cache_configs[0].num_blocks

204 - for cfg in kv_cache_configs

205 - ])

206 - num_gpu_blocks = kv_cache_configs[0].num_blocks

201 + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(

202 + kv_cache_configs)

203 + num_gpu_blocks = scheduler_kv_cache_config.num_blocks

204 num_cpu_blocks = 0

208 - scheduler_kv_cache_config = kv_cache_configs[0]

205

206 # Initialize kv cache and warmup the execution

207 self.model_executor.initialize_from_config(kv_cache_configs)

vllm/v1/kv_cache_interface.py

+70

@@ -234,6 +234,76 @@ class CrossAttentionSpec(AttentionSpec):

234 return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes

235

236

237 +@dataclass(frozen=True)

238 +class UniformTypeKVCacheSpecs(KVCacheSpec):

239 + """

240 + A KV cache spec for multiple layers with the same type of attention. Here,

241 + same types means always need the same number of token slots. For example,

242 + sliding window attentions with different window sizes are not the same type

243 + and should not be merged into one UniformTypeKVCacheSpecs.

244 + """

245 + kv_cache_specs: dict[str, KVCacheSpec]

246 +

247 + @property

248 + def page_size_bytes(self) -> int:

249 + return sum(spec.page_size_bytes

250 + for spec in self.kv_cache_specs.values())

251 +

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

253 + max_num_pages = max(

254 + cdiv(spec.max_memory_usage_bytes(vllm_config),

255 + spec.page_size_bytes)

256 + for spec in self.kv_cache_specs.values())

257 + return max_num_pages * self.page_size_bytes

258 +

259 + @classmethod

260 + def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:

261 + """

262 + Whether all layers have the same type of KV cache spec.

263 + """

264 + block_sizes = set(spec.block_size for spec in kv_cache_specs.values())

265 + if len(block_sizes) > 1:

266 + # Different block sizes, not uniform.

267 + return False

268 + one_spec = next(iter(kv_cache_specs.values()))

269 + if isinstance(one_spec, (FullAttentionSpec, CrossAttentionSpec)):

270 + return all(

271 + isinstance(spec, type(one_spec))

272 + for spec in kv_cache_specs.values())

273 + elif isinstance(one_spec, SlidingWindowSpec):

274 + return all(

275 + isinstance(spec, SlidingWindowSpec)

276 + and spec.sliding_window == one_spec.sliding_window

277 + for spec in kv_cache_specs.values())

278 + elif isinstance(one_spec, ChunkedLocalAttentionSpec):

279 + return all(

280 + isinstance(spec, ChunkedLocalAttentionSpec)

281 + and spec.attention_chunk_size == one_spec.attention_chunk_size

282 + for spec in kv_cache_specs.values())

283 + elif isinstance(one_spec, MambaSpec):

284 + return all(

285 + isinstance(spec, MambaSpec) and spec.num_speculative_blocks ==

286 + one_spec.num_speculative_blocks

287 + for spec in kv_cache_specs.values())

288 + else:

289 + # NOTE(Chen): Please add new branches for new KV cache spec types.

290 + raise NotImplementedError(

291 + f"Unsupported KV cache spec type: {type(one_spec)}")

292 +

293 + @classmethod

294 + def from_specs(cls, kv_cache_specs: dict[str,

295 + KVCacheSpec]) -> Optional[Self]:

296 + """

297 + Return a SameTypeKVCacheSpecs object if all layers have the same type

298 + of KV cache spec. Return None if not.

299 + """

300 + if cls.is_uniform_type(kv_cache_specs):

301 + block_size = next(iter(kv_cache_specs.values())).block_size

302 + return cls(block_size=block_size, kv_cache_specs=kv_cache_specs)

303 + else:

304 + return None

305 +

306 +

307 @dataclass

308 class KVCacheTensor:

309 """

vllm/v1/worker/gpu_model_runner.py

+36 -29

@@ -8,7 +8,7 @@ from collections import defaultdict

8 from collections.abc import Iterator

9 from contextlib import contextmanager

10 from copy import deepcopy

11 -from typing import TYPE_CHECKING, Any, Optional, Union, cast

11 +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union, cast

13 import numpy as np

14 import torch

@@ -74,7 +74,8 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,

74 EncoderOnlyAttentionSpec,

75 FullAttentionSpec, KVCacheConfig,

76 KVCacheGroupSpec, KVCacheSpec,

77 - MambaSpec, SlidingWindowSpec)

77 + MambaSpec, SlidingWindowSpec,

78 + UniformTypeKVCacheSpecs)

79 # yapf: enable

80 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,

81 DraftTokenIds, LogprobsLists, LogprobsTensors,

@@ -1187,7 +1188,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1188 common_prefix_len = self._compute_cascade_attn_prefix_len(

1189 num_scheduled_tokens,

1190 num_common_prefix_blocks,

1190 - kv_cache_group_spec.kv_cache_spec,

1191 + attn_group.kv_cache_spec,

1192 builder,

1193 )

1194

@@ -3453,12 +3454,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3454 assert len(self.attn_groups) == 0, \

3455 "Attention backends are already initialized"

3456

3456 - def get_attn_backends_for_layers(

3457 - layer_names: list[str]

3458 - ) -> dict[type[AttentionBackend], list[str]]:

3459 - layers = get_layers_from_vllm_config(self.vllm_config,

3460 - AttentionLayerBase,

3461 - layer_names)

3457 + class AttentionGroupKey(NamedTuple):

3458 + attn_backend: type[AttentionBackend]

3459 + kv_cache_spec: KVCacheSpec

3460 +

3461 + def get_attn_backends_for_group(

3462 + kv_cache_group_spec: KVCacheGroupSpec,

3463 + ) -> dict[AttentionGroupKey, list[str]]:

3464 + layers = get_layers_from_vllm_config(

3465 + self.vllm_config, AttentionLayerBase,

3466 + kv_cache_group_spec.layer_names)

3467 attn_backends = {}

3468 attn_backend_layers = defaultdict(list)

3469 # Dedupe based on full class name; this is a bit safer than

@@ -3466,7 +3471,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3471 # attention backend subclasses (e.g. ChunkedLocalAttention) unless

3472 # they are cached correctly, there will be different objects per

3473 # layer.

3469 - for layer_name in layer_names:

3474 + for layer_name in kv_cache_group_spec.layer_names:

3475 attn_backend = layers[layer_name].get_attn_backend()

3476

3477 if layer_name in self.kv_sharing_fast_prefill_eligible_layers:

@@ -3475,8 +3480,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3480 attn_backend,

3481 )

3482

3478 - key = attn_backend.full_cls_name()

3479 - attn_backends[key] = attn_backend

3483 + full_cls_name = attn_backend.full_cls_name()

3484 + layer_kv_cache_spec = kv_cache_group_spec.kv_cache_spec

3485 + if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):

3486 + layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[

3487 + layer_name]

3488 + key = (full_cls_name, layer_kv_cache_spec)

3489 + attn_backends[key] = AttentionGroupKey(attn_backend,

3490 + layer_kv_cache_spec)

3491 attn_backend_layers[key].append(layer_name)

3492 return {

3493 attn_backends[k]: v

@@ -3484,11 +3495,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3495 }

3496

3497 def create_attn_groups(

3487 - attn_backends_map: dict[AttentionBackend, list[str]],

3488 - kv_cache_spec: KVCacheSpec,

3498 + attn_backends_map: dict[AttentionGroupKey, list[str]],

3499 ) -> list[AttentionGroup]:

3500 attn_groups: list[AttentionGroup] = []

3491 - for attn_backend, layer_names in attn_backends_map.items():

3501 + for (attn_backend,

3502 + kv_cache_spec), layer_names in attn_backends_map.items():

3503 attn_metadata_builders = []

3504 attn_metadata_builders.append(attn_backend.get_builder_cls()(

3505 kv_cache_spec,

@@ -3506,16 +3517,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3517 ))

3518 attn_group = AttentionGroup(attn_backend,

3519 attn_metadata_builders,

3509 - layer_names)

3520 + layer_names, kv_cache_spec)

3521 attn_groups.append(attn_group)

3522 return attn_groups

3523

3524 for kv_cache_group_spec in kv_cache_config.kv_cache_groups:

3514 - kv_cache_spec = kv_cache_group_spec.kv_cache_spec

3515 - attn_backends = get_attn_backends_for_layers(

3516 - kv_cache_group_spec.layer_names)

3517 - self.attn_groups.append(

3518 - create_attn_groups(attn_backends, kv_cache_spec))

3525 + attn_backends = get_attn_backends_for_group(kv_cache_group_spec)

3526 + self.attn_groups.append(create_attn_groups(attn_backends))

3527

3528 # Calculate reorder batch threshold (if needed)

3529 self.calculate_reorder_batch_threshold()

@@ -3680,14 +3688,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3688 def _attn_group_iterator(self) -> Iterator[AttentionGroup]:

3689 return itertools.chain.from_iterable(self.attn_groups)

3690

3683 - def _kv_cache_spec_attn_group_iterator(

3684 - self) -> Iterator[tuple[KVCacheSpec, AttentionGroup]]:

3691 + def _kv_cache_spec_attn_group_iterator(self) -> Iterator[AttentionGroup]:

3692 if not self.kv_cache_config.kv_cache_groups:

3693 return

3687 - for kv_cache_spec_id, attn_groups in enumerate(self.attn_groups):

3688 - for attn_group in attn_groups:

3689 - yield self.kv_cache_config.kv_cache_groups[

3690 - kv_cache_spec_id].kv_cache_spec, attn_group

3694 + for attn_groups in self.attn_groups:

3695 + yield from attn_groups

3696

3697 def _reshape_kv_cache_tensors(

3698 self,

@@ -3707,7 +3712,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3712 """

3713 kv_caches: dict[str, torch.Tensor] = {}

3714 has_attn, has_mamba = False, False

3710 - for kv_cache_spec, group in self._kv_cache_spec_attn_group_iterator():

3715 + for group in self._kv_cache_spec_attn_group_iterator():

3716 + kv_cache_spec = group.kv_cache_spec

3717 attn_backend = group.backend

3718 for layer_name in group.layer_names:

3719 if layer_name in self.runner_only_attn_layers:

@@ -3787,7 +3793,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3793 kv_caches: The KV cache buffer of each layer.

3794 """

3795

3790 - for kv_cache_spec, group in self._kv_cache_spec_attn_group_iterator():

3796 + for group in self._kv_cache_spec_attn_group_iterator():

3797 + kv_cache_spec = group.kv_cache_spec

3798 for layer_name in group.layer_names:

3799 kv_cache = kv_caches[layer_name]

3800 if (isinstance(kv_cache_spec, AttentionSpec)

vllm/v1/worker/utils.py

+2 -1

@@ -15,7 +15,7 @@ from vllm.multimodal.registry import MultiModalRegistry

15 from vllm.platforms import current_platform

16 from vllm.v1.attention.backends.utils import AttentionMetadataBuilder

17 from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget

18 -from vllm.v1.kv_cache_interface import KVCacheGroupSpec

18 +from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec

20 if TYPE_CHECKING:

21 from vllm.attention.layer import Attention

@@ -132,6 +132,7 @@ class AttentionGroup:

132 backend: type[AttentionBackend]

133 metadata_builders: list[AttentionMetadataBuilder]

134 layer_names: list[str]

135 + kv_cache_spec: KVCacheSpec

136

137 def get_metadata_builder(self,

138 ubatch_id: Optional[int] = None