@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 93a00d7dde
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 15, 18:30 UTC

[v1] Refactor KVCacheConfig (#14079)

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

Chen Zhang committed Mar 21, 2025 at 19:56 UTC 93a00d7ddec29371efb4764d4c55065eca4c7746
10 files changed +320 -112
tests/v1/core/test_kv_cache_utils.py
+109 -1
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import pytest
4 +import torch
5
6 from vllm.multimodal.inputs import MultiModalKwargs
7 from vllm.sampling_params import SamplingParams
@@ -8,7 +9,10 @@ from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
9 KVCacheBlock, PrefixCachingMetrics,
10 generate_block_hash_extra_keys,
11 hash_block_tokens,
11 - hash_request_tokens)
12 + hash_request_tokens,
13 + unify_kv_cache_configs)
14 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
15 + KVCacheGroupSpec, KVCacheTensor)
16 from vllm.v1.metrics.stats import PrefixCacheStats
17 from vllm.v1.request import Request
18
@@ -314,3 +318,107 @@ def test_metrics():
318 assert metrics.aggregated_query_total == 0
319 assert metrics.aggregated_query_hit == 0
320 assert not metrics.query_queue
321 +
322 +
323 +def test_unify_kv_cache_configs():
324 +
325 + def new_kv_cache_spec(block_size=16,
326 + num_kv_heads=2,
327 + head_size=64,
328 + dtype=torch.float32,
329 + use_mla=False):
330 + return FullAttentionSpec(block_size=block_size,
331 + num_kv_heads=num_kv_heads,
332 + head_size=head_size,
333 + dtype=dtype,
334 + use_mla=use_mla)
335 +
336 + same_kv_cache_config = [
337 + KVCacheConfig(
338 + num_blocks=10,
339 + tensors={
340 + "layer1": KVCacheTensor(100),
341 + "layer2": KVCacheTensor(100),
342 + },
343 + kv_cache_groups=[
344 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
345 + KVCacheGroupSpec(["layer2"],
346 + new_kv_cache_spec(num_kv_heads=4)),
347 + ],
348 + ),
349 + KVCacheConfig(
350 + num_blocks=20,
351 + tensors={
352 + "layer1": KVCacheTensor(100),
353 + "layer2": KVCacheTensor(100),
354 + },
355 + kv_cache_groups=[
356 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
357 + KVCacheGroupSpec(["layer2"],
358 + new_kv_cache_spec(num_kv_heads=4)),
359 + ],
360 + ),
361 + ]
362 + unify_kv_cache_configs(same_kv_cache_config)
363 + assert same_kv_cache_config[0].num_blocks == 10
364 + assert same_kv_cache_config[1].num_blocks == 10
365 +
366 + need_sort_kv_cache_config = [
367 + KVCacheConfig(
368 + num_blocks=10,
369 + tensors={
370 + "layer1": KVCacheTensor(100),
371 + "layer2": KVCacheTensor(100),
372 + },
373 + kv_cache_groups=[
374 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
375 + KVCacheGroupSpec(["layer2"],
376 + new_kv_cache_spec(num_kv_heads=4)),
377 + ],
378 + ),
379 + KVCacheConfig(
380 + num_blocks=20,
381 + tensors={
382 + "layer1": KVCacheTensor(100),
383 + "layer2": KVCacheTensor(100),
384 + },
385 + kv_cache_groups=[
386 + KVCacheGroupSpec(["layer2"],
387 + new_kv_cache_spec(num_kv_heads=4)),
388 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
389 + ],
390 + ),
391 + ]
392 +
393 + unify_kv_cache_configs(need_sort_kv_cache_config)
394 + assert need_sort_kv_cache_config[0].num_blocks == 10
395 + assert need_sort_kv_cache_config[1].num_blocks == 10
396 +
397 + diff_kv_cache_config = [
398 + KVCacheConfig(
399 + num_blocks=10,
400 + tensors={
401 + "layer1": KVCacheTensor(100),
402 + "layer2": KVCacheTensor(100),
403 + },
404 + kv_cache_groups=[
405 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
406 + KVCacheGroupSpec(["layer2"],
407 + new_kv_cache_spec(num_kv_heads=4)),
408 + ],
409 + ),
410 + KVCacheConfig(
411 + num_blocks=20,
412 + tensors={
413 + "layer1": KVCacheTensor(100),
414 + "layer2": KVCacheTensor(100),
415 + },
416 + kv_cache_groups=[
417 + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
418 + KVCacheGroupSpec(["layer2"],
419 + new_kv_cache_spec(num_kv_heads=8)),
420 + ],
421 + ),
422 + ]
423 + with pytest.raises(AssertionError):
424 + unify_kv_cache_configs(diff_kv_cache_config)
vllm/v1/core/kv_cache_utils.py
+95 -35
@@ -7,8 +7,8 @@ from typing import Any, NamedTuple, Optional
7
8 from vllm.config import VllmConfig
9 from vllm.logger import init_logger
10 -from vllm.v1.kv_cache_interface import (KVCacheConfig, KVCacheSpec,
11 - KVCacheTensor)
10 +from vllm.v1.kv_cache_interface import (KVCacheConfig, KVCacheGroupSpec,
11 + KVCacheSpec, KVCacheTensor)
12 from vllm.v1.metrics.stats import PrefixCacheStats
13 from vllm.v1.request import Request
14
@@ -449,7 +449,7 @@ def hash_request_tokens(block_size: int,
449
450
451 def check_enough_kv_cache_memory(vllm_config: VllmConfig,
452 - kv_cache_spec: KVCacheSpec,
452 + kv_cache_spec: dict[str, KVCacheSpec],
453 available_memory: int):
454 """
455 Checks whether `available_memory` is enough for the KV cache to hold at
@@ -457,7 +457,7 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,
457
458 Args:
459 vllm_config: The global VllmConfig
460 - kv_cache_spec: The kv cache spec of the model
460 + kv_cache_spec: The kv cache spec of each attention layer in the model
461 available_memory: Memory available for KV cache in bytes.
462
463 Raises:
@@ -484,12 +484,43 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,
484 f"`max_model_len` when initializing the engine.")
485
486
487 -def is_kv_cache_type_uniform(kv_cache_spec: KVCacheSpec) -> bool:
487 +def create_kv_cache_group_specs(
488 + kv_cache_spec: dict[str, KVCacheSpec],
489 + grouped_layer_names: list[list[str]]) -> list[KVCacheGroupSpec]:
490 + """
491 + Create KVCacheGroupSpec object for each kv cache group layer.
492 + The layers in the same group should share the same
493 + KVCacheSpec.
494 +
495 + Args:
496 + kv_cache_spec:
497 + A mapping from each layer name to its corresponding KVCacheSpec.
498 + grouped_layer_names:
499 + A list of kv cache groups, where each element is a list of layer
500 + names that belong to the same group and should share the same
501 + KVCacheSpec.
502 + Returns:
503 + A list of KVCacheGroupSpec objects, one for each group.
504 + """
505 + kv_cache_groups = []
506 + for layer_names_one_group in grouped_layer_names:
507 + layer_spec = kv_cache_spec[layer_names_one_group[0]]
508 + assert all(
509 + kv_cache_spec[layer_name] == layer_spec
510 + for layer_name in layer_names_one_group[1:]), (
511 + "All layers in the same KV cache group must share the same "
512 + "KVCacheSpec.")
513 + kv_cache_groups.append(
514 + KVCacheGroupSpec(layer_names_one_group, layer_spec))
515 + return kv_cache_groups
516 +
517 +
518 +def is_kv_cache_type_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
519 """
520 Whether all layers in the given KVCacheSpec have the same type of KV cache.
521
522 Args:
492 - kv_cache_spec: The KVCacheSpec of the model
523 + kv_cache_spec: The kv cache spec of each attention layer in the model
524
525 Returns:
526 True if all layers have the same type, False otherwise.
@@ -500,18 +531,16 @@ def is_kv_cache_type_uniform(kv_cache_spec: KVCacheSpec) -> bool:
531
532
533 def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
503 - kv_cache_spec: KVCacheSpec,
504 - available_memory: int,
505 - num_layers: int) -> KVCacheConfig:
534 + kv_cache_spec: dict[str, KVCacheSpec],
535 + available_memory: int) -> KVCacheConfig:
536 """
537 Generates the KV cache configuration for a model with one type of KV cache.
538 Divide the available memory equally among all layers.
539
540 Args:
541 vllm_config: The global VllmConfig
512 - kv_cache_spec: The kv cache spec of the model
542 + kv_cache_spec: The kv cache spec of each attention layer in the model
543 available_memory: Memory available for KV cache in bytes.
514 - num_layers: The number of layers in the model.
544
545 Returns:
546 The generated KVCacheConfig
@@ -521,7 +550,7 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
550 assert len(page_sizes) == 1
551 page_size = page_sizes.pop()
552
524 - num_blocks = int(available_memory // page_size // num_layers)
553 + num_blocks = int(available_memory // page_size // len(kv_cache_spec))
554 num_blocks = max(num_blocks, 0)
555
556 if vllm_config.cache_config.num_gpu_blocks_override is not None:
@@ -541,6 +570,9 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
570 max_model_len_str, max_concurrency)
571
572 per_layer_size = page_size * num_blocks
573 + # All layers have the same KV cache spec, so we create one kv cache group
574 + # for all layers.
575 + grouped_layer_names = [list(kv_cache_spec.keys())]
576
577 kv_cache_config = KVCacheConfig(
578 num_blocks=num_blocks,
@@ -548,41 +580,69 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
580 layer_name: KVCacheTensor(size=per_layer_size)
581 for layer_name in kv_cache_spec
582 },
551 - groups=[[layer_name for layer_name in kv_cache_spec]],
552 - kv_cache_spec=kv_cache_spec)
583 + kv_cache_groups=create_kv_cache_group_specs(kv_cache_spec,
584 + grouped_layer_names),
585 + )
586 return kv_cache_config
587
588
556 -def get_kv_cache_configs(vllm_config: VllmConfig,
557 - kv_cache_specs: list[KVCacheSpec],
558 - available_memory: int) -> list[KVCacheConfig]:
589 +def get_kv_cache_config(vllm_config: VllmConfig,
590 + kv_cache_spec: dict[str, KVCacheSpec],
591 + available_memory: int) -> KVCacheConfig:
592 """
593 Generates the KV cache configuration for a model
594 TODO: support hybrid models with more than one type of KV cache.
595
596 Args:
597 vllm_config: The global VllmConfig
565 - kv_cache_specs: The kv cache specs of the model
598 + kv_cache_spec: The kv cache spec of each attention layer in the model
599 available_memory: Memory available for KV cache in bytes.
600
601 Returns:
602 The generated KVCacheConfigs
603 """
571 - # Use the max number of layers to conservatively determine
572 - # the number of blocks.
573 - num_layers = max(len(kv_cache_spec) for kv_cache_spec in kv_cache_specs)
574 - kv_cache_configs = []
575 - for kv_cache_spec in kv_cache_specs:
576 - check_enough_kv_cache_memory(vllm_config, kv_cache_spec,
577 - available_memory)
578 - if is_kv_cache_type_uniform(kv_cache_spec):
579 - # KV cache of all layers are the same, which is true for
580 - # most models. Allocate the same amount of memory for
581 - # each layer.
582 - kv_cache_configs.append(
583 - _get_kv_cache_config_uniform_type(vllm_config, kv_cache_spec,
584 - available_memory,
585 - num_layers))
586 - else:
587 - raise NotImplementedError
604 + check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
605 + if is_kv_cache_type_uniform(kv_cache_spec):
606 + # KV cache of all layers are the same, which is true for
607 + # most models. Allocate the same amount of memory for
608 + # each layer.
609 + return _get_kv_cache_config_uniform_type(vllm_config, kv_cache_spec,
610 + available_memory)
611 +
612 + raise NotImplementedError
613 +
614 +
615 +def unify_kv_cache_configs(kv_cache_configs: list[KVCacheConfig]):
616 + """
617 + Make the KV cache configurations for each worker consistent, so that all
618 + workers can be controlled by the same KVCacheManager.
619 + This function verifies that the layer group of each worker are the same,
620 + and changes the num_blocks of each worker to the smallest among all workers.
621 +
622 + Args:
623 + kv_cache_configs: The KV cache configurations for each worker. Will be
624 + in-place modified to make them consistent.
625 + """
626 +
627 + # Sort the kv cache groups by the type_id of their KV cache spec.
628 + # This can avoid the inconsistency caused by the order of groups.
629 + for kv_cache_config in kv_cache_configs:
630 + kv_cache_config.kv_cache_groups.sort(
631 + key=lambda x: x.kv_cache_spec.type_id)
632 +
633 + # Verify that the groups of each rank are the same.
634 + for kv_cache_config in kv_cache_configs[1:]:
635 + for group_rank_0, group_rank_i in zip(
636 + kv_cache_configs[0].kv_cache_groups,
637 + kv_cache_config.kv_cache_groups):
638 + assert group_rank_0.kv_cache_spec == group_rank_i.kv_cache_spec
639 +
640 + # Change the num_blocks of each rank to the smallest among all ranks. We
641 + # do not need to shrink the tensor size because it is valid to only use the
642 + # first `num_blocks` blocks of the tensor.
643 + min_num_blocks = min(kv_cache_config.num_blocks
644 + for kv_cache_config in kv_cache_configs)
645 + for kv_cache_config in kv_cache_configs:
646 + kv_cache_config.num_blocks = min_num_blocks
647 +
648 return kv_cache_configs
vllm/v1/engine/core.py
+22 -9
@@ -21,7 +21,8 @@ from vllm.transformers_utils.config import (
21 maybe_register_config_serialize_by_value)
22 from vllm.utils import (get_exception_traceback, resolve_obj_by_qualname,
23 zmq_socket_ctx)
24 -from vllm.v1.core.kv_cache_utils import get_kv_cache_configs
24 +from vllm.v1.core.kv_cache_utils import (get_kv_cache_config,
25 + unify_kv_cache_configs)
26 from vllm.v1.core.sched.output import SchedulerOutput
27 from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler
28 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
@@ -120,15 +121,27 @@ class EngineCore:
121 # memory can be allocated for kv cache.
122 available_gpu_memory = self.model_executor.determine_available_memory()
123
124 + assert len(kv_cache_specs) == len(available_gpu_memory)
125 # Get the kv cache tensor size
124 - kv_cache_configs = get_kv_cache_configs(vllm_config, kv_cache_specs,
125 - available_gpu_memory)
126 - num_gpu_blocks_set = set(config.num_blocks
127 - for config in kv_cache_configs)
128 - assert len(num_gpu_blocks_set) == 1, (
129 - f"num_gpu_blocks need to be the same across workers, "
130 - f"but they are different: {num_gpu_blocks_set}")
131 - num_gpu_blocks = num_gpu_blocks_set.pop()
126 + kv_cache_configs = [
127 + get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
128 + available_gpu_memory_one_worker)
129 + for kv_cache_spec_one_worker, available_gpu_memory_one_worker in
130 + zip(kv_cache_specs, available_gpu_memory)
131 + ]
132 +
133 + # Since we use a shared centralized controller, we need the
134 + # `kv_cache_config` to be consistent across all workers to make sure
135 + # all the memory operators can be applied to all workers.
136 + unify_kv_cache_configs(kv_cache_configs)
137 +
138 + # All workers have the same kv_cache_config except layer names, so use
139 + # an arbitrary one to get the number of blocks.
140 + assert all([
141 + cfg.num_blocks == kv_cache_configs[0].num_blocks
142 + for cfg in kv_cache_configs
143 + ])
144 + num_gpu_blocks = kv_cache_configs[0].num_blocks
145 num_cpu_blocks = 0
146
147 # Initialize kv cache and warmup the execution
vllm/v1/executor/abstract.py
+5 -8
@@ -62,14 +62,11 @@ class Executor(ExecutorBase):
62 args=(kv_cache_configs, ))
63 self.collective_rpc("compile_or_warm_up_model")
64
65 - def determine_available_memory(self) -> int: # in bytes
65 + def determine_available_memory(self) -> list[int]: # in bytes
66 output = self.collective_rpc("determine_available_memory")
67 - # Since we use a shared centralized controller, we take the minimum
68 - # memory size across all workers to make sure all the memory
69 - # operators can be applied to all workers.
70 - return min(output)
67 + return output
68
72 - def get_kv_cache_specs(self) -> list[KVCacheSpec]:
69 + def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
70 output = self.collective_rpc("get_kv_cache_spec")
71 return output
72
@@ -95,7 +92,7 @@ class UniProcExecutor(UniProcExecutorV0, Executor):
92
93 class ExecutorWithExternalLauncher(ExecutorWithExternalLauncherV0, Executor):
94
98 - def determine_available_memory(self) -> int: # in bytes
95 + def determine_available_memory(self) -> list[int]: # in bytes
96 # same as determine_num_available_blocks in v0,
97 # we need to get the min across all ranks.
98 memory = super().determine_available_memory()
@@ -103,4 +100,4 @@ class ExecutorWithExternalLauncher(ExecutorWithExternalLauncherV0, Executor):
100 cpu_group = get_world_group().cpu_group
101 memory_tensor = torch.tensor([memory], device="cpu", dtype=torch.int64)
102 dist.all_reduce(memory_tensor, group=cpu_group, op=dist.ReduceOp.MIN)
106 - return memory_tensor.item()
103 + return [memory_tensor.item()]
vllm/v1/kv_cache_interface.py
+34 -18
@@ -11,7 +11,7 @@ logger = init_logger(__name__)
11
12
13 @dataclass
14 -class KVCacheSpecBase:
14 +class KVCacheSpec:
15 """
16 A base class for specifying the KV cache format of one layer.
17 """
@@ -55,7 +55,7 @@ class KVCacheSpecBase:
55
56
57 @dataclass
58 -class FullAttentionSpec(KVCacheSpecBase):
58 +class FullAttentionSpec(KVCacheSpec):
59 num_kv_heads: int
60 head_size: int
61 dtype: torch.dtype
@@ -76,9 +76,6 @@ class FullAttentionSpec(KVCacheSpecBase):
76 return cdiv(num_tokens, self.block_size) * self.page_size_bytes
77
78
79 -KVCacheSpec = dict[str, KVCacheSpecBase]
80 -
81 -
79 @dataclass
80 class KVCacheTensor:
81 """
@@ -89,6 +86,18 @@ class KVCacheTensor:
86 size: int # The size of KV cache Tensor in bytes
87
88
89 +@dataclass
90 +class KVCacheGroupSpec:
91 + """
92 + Represents a group of model layers that share the same KV cache block table.
93 + These layers are regarded as one layer in the KV cache manager.
94 + """
95 + # The names of model layers in this group
96 + layer_names: list[str]
97 + # The KV cache spec of this manager layer
98 + kv_cache_spec: KVCacheSpec
99 +
100 +
101 @dataclass
102 class KVCacheConfig:
103 """
@@ -99,17 +108,24 @@ class KVCacheConfig:
108 """layer_name -> how to initialize KV cache for that layer"""
109 tensors: dict[str, KVCacheTensor]
110 """
102 - A list of kv-cache groups. Each group includes a set of layers with
103 - the same kv-cache spec, and the total page_size of layers inside a group
104 - is same across all groups (as the KVCacheManager only supports allocating
105 - pages of the same size). For example:
106 - 1. A model only uses full attention: one group with all layers in the model.
107 - 2. (not implemented yet) A model with the same number of full attention
108 - layers and sliding window attention layers: two groups, one for full
109 - attention layers and one for sliding window attention layers.
110 - 3. (not implemented yet) A model with 2 full attention layers and 4 sliding
111 - window attention layers: three groups, (full * 2), (sw * 2), (sw * 2).
111 + The kv cache groups of the model.
112 + The layers in the models are repeated with some patterns, e.g., a model
113 + with 10 full attention layers and 20 sliding window attention layers can be
114 + regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
115 + The KVCacheManager allocates different block tables for each of the 3 layers
116 + in the pattern, and repeats each of them 10 times to generate the
117 + block_table for the 30 layers in the model.
118 + Therefore, we can group the layers in the model into 3 groups, each of which
119 + contains 10 layers in the model.
120 + The KVCacheManager allocates the block_table for each group based on its
121 + kv_cache spec, and the model runner applies the block table to each layer
122 + in the group.
123 + For example:
124 + 1. A model only uses full attention. The pattern is
125 + (num_hidden_layers * full), so there is only one group and the block table
126 + is shared by all layers.
127 + 2. (WIP) A model with 10 full attention layers and 20 sliding window
128 + attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
129 + there are 3 groups, each of which represents 10 layers in the model.
130 """
113 - groups: list[list[str]]
114 - """the KVCacheSpec of the model"""
115 - kv_cache_spec: KVCacheSpec
131 + kv_cache_groups: list[KVCacheGroupSpec]
vllm/v1/worker/gpu_model_runner.py
+29 -17
@@ -1510,34 +1510,46 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1510 kv_cache_config: Configuration for the KV cache, including the KV
1511 cache size of each layer
1512 """
1513 - if len(kv_cache_config.groups) > 1:
1513 + if len(kv_cache_config.kv_cache_groups) > 1:
1514 raise NotImplementedError(
1515 "Hybrid models with more than one KV cache type are not "
1516 "supported yet.")
1517
1518 kv_caches: dict[str, torch.Tensor] = {}
1519
1520 - for layer_name, layer_spec in kv_cache_config.kv_cache_spec.items():
1521 - tensor_config = kv_cache_config.tensors[layer_name]
1522 - assert tensor_config.size % layer_spec.page_size_bytes == 0
1523 - num_blocks = tensor_config.size // layer_spec.page_size_bytes
1524 - if isinstance(layer_spec, FullAttentionSpec):
1525 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(
1526 - num_blocks, layer_spec.block_size, layer_spec.num_kv_heads,
1527 - layer_spec.head_size)
1528 - dtype = layer_spec.dtype
1529 - kv_caches[layer_name] = torch.zeros(kv_cache_shape,
1530 - dtype=dtype,
1531 - device=self.device)
1532 - else:
1533 - raise NotImplementedError
1520 + for kv_cache_group in kv_cache_config.kv_cache_groups:
1521 + kv_cache_spec = kv_cache_group.kv_cache_spec
1522 + for layer_name in kv_cache_group.layer_names:
1523 + tensor_config = kv_cache_config.tensors[layer_name]
1524 + assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
1525 + num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
1526 + # `num_blocks` is the number of blocks the model runner can use.
1527 + # `kv_cache_config.num_blocks` is the number of blocks that
1528 + # KVCacheManager may allocate.
1529 + # Since different GPUs may have different number of layers and
1530 + # different memory capacities, `num_blocks` can be different on
1531 + # different GPUs, and `kv_cache_config.num_blocks` is set to
1532 + # the min of all `num_blocks`. Verify it here.
1533 + assert num_blocks >= kv_cache_config.num_blocks
1534 + if isinstance(kv_cache_spec, FullAttentionSpec):
1535 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(
1536 + num_blocks, kv_cache_spec.block_size,
1537 + kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
1538 + dtype = kv_cache_spec.dtype
1539 + kv_caches[layer_name] = torch.zeros(kv_cache_shape,
1540 + dtype=dtype,
1541 + device=self.device)
1542 + else:
1543 + # TODO: add new branches when introducing more types of
1544 + # KV cache specs.
1545 + raise ValueError("Unknown KV cache spec type.")
1546
1547 bind_kv_cache(
1548 kv_caches,
1549 self.vllm_config.compilation_config.static_forward_context,
1550 self.kv_caches)
1551
1540 - def get_kv_cache_spec(self) -> KVCacheSpec:
1552 + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
1553 """
1554 Generates the KVCacheSpec by parsing the kv cache format from each
1555 Attention module in the static forward context.
@@ -1549,7 +1561,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1561 forward_ctx = self.vllm_config.compilation_config.static_forward_context
1562 block_size = self.vllm_config.cache_config.block_size
1563 use_mla = self.vllm_config.model_config.use_mla
1552 - kv_cache_spec: KVCacheSpec = {}
1564 + kv_cache_spec: dict[str, KVCacheSpec] = {}
1565 for layer_name, attn_module in forward_ctx.items():
1566 if isinstance(attn_module, FusedMoE):
1567 continue
vllm/v1/worker/gpu_worker.py
+1 -1
@@ -185,7 +185,7 @@ class Worker(WorkerBase):
185
186 return int(available_kv_cache_memory)
187
188 - def get_kv_cache_spec(self) -> KVCacheSpec:
188 + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
189 return self.model_runner.get_kv_cache_spec()
190
191 def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
vllm/v1/worker/tpu_model_runner.py
+23 -21
@@ -309,7 +309,7 @@ class TPUModelRunner:
309 assert self.model is not None
310 return self.model
311
312 - def get_kv_cache_spec(self) -> KVCacheSpec:
312 + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
313 """
314 Generates the KVCacheSpec by parsing the kv cache format from each
315 Attention module in the static forward context.
@@ -320,7 +320,7 @@ class TPUModelRunner:
320
321 forward_ctx = self.vllm_config.compilation_config.static_forward_context
322 block_size = self.vllm_config.cache_config.block_size
323 - kv_cache_spec: KVCacheSpec = {}
323 + kv_cache_spec: dict[str, KVCacheSpec] = {}
324 for layer_name, attn_module in forward_ctx.items():
325 # TODO: Support other attention modules, e.g., sliding window,
326 # cross-attention, MLA.
@@ -837,31 +837,33 @@ class TPUModelRunner:
837 kv_cache_config: Configuration for the KV cache, including the KV
838 cache size of each layer
839 """
840 - if len(kv_cache_config.groups) > 1:
840 + if len(kv_cache_config.kv_cache_groups) > 1:
841 raise NotImplementedError(
842 "Hybrid models with more than one KV cache type are not "
843 "supported yet.")
844
845 kv_caches: dict[str, torch.Tensor] = {}
846
847 - for layer_name, layer_spec in kv_cache_config.kv_cache_spec.items():
848 - tensor_config = kv_cache_config.tensors[layer_name]
849 - assert tensor_config.size % layer_spec.page_size_bytes == 0
850 - num_blocks = tensor_config.size // layer_spec.page_size_bytes
851 - if isinstance(layer_spec, FullAttentionSpec):
852 - kv_cache_shape = PallasAttentionBackend.get_kv_cache_shape(
853 - num_blocks, layer_spec.block_size, layer_spec.num_kv_heads,
854 - layer_spec.head_size)
855 - dtype = layer_spec.dtype
856 -
857 - tpu_k_cache = torch.zeros(kv_cache_shape,
858 - dtype=dtype,
859 - device=self.device)
860 - tpu_v_cache = torch.zeros_like(tpu_k_cache)
861 -
862 - kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
863 - else:
864 - raise NotImplementedError
847 + for kv_cache_group in kv_cache_config.kv_cache_groups:
848 + kv_cache_spec = kv_cache_group.kv_cache_spec
849 + for layer_name in kv_cache_group.layer_names:
850 + tensor_config = kv_cache_config.tensors[layer_name]
851 + assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
852 + num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
853 + if isinstance(kv_cache_spec, FullAttentionSpec):
854 + kv_cache_shape = PallasAttentionBackend.get_kv_cache_shape(
855 + num_blocks, kv_cache_spec.block_size,
856 + kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
857 + dtype = kv_cache_spec.dtype
858 +
859 + tpu_k_cache = torch.zeros(kv_cache_shape,
860 + dtype=dtype,
861 + device=self.device)
862 + tpu_v_cache = torch.zeros_like(tpu_k_cache)
863 +
864 + kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
865 + else:
866 + raise NotImplementedError
867
868 bind_kv_cache(
869 kv_caches,
vllm/v1/worker/tpu_worker.py
+1 -1
@@ -189,7 +189,7 @@ class TPUWorker:
189 def get_model(self) -> nn.Module:
190 return self.model_runner.get_model()
191
192 - def get_kv_cache_spec(self) -> KVCacheSpec:
192 + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
193 return self.model_runner.get_kv_cache_spec()
194
195 def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
vllm/v1/worker/worker_base.py
+1 -1
@@ -51,7 +51,7 @@ class WorkerBase(WorkerBaseV0):
51 self.device: Optional[torch.device] = None
52 self.model_runner: Optional[nn.Module] = None
53
54 - def get_kv_cache_spec(self) -> KVCacheSpec:
54 + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
55 """Get specifications for KV cache implementation."""
56 raise NotImplementedError
57