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
[Attention] Refactor AttentionMetadata Preparation for Encoder-only Models (#23154)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Aug 21, 2025 at 22:05 UTC
17373dcd93ca60554d72cef4e159e70abbfd15af
12 files changed
+227
-215
tests/v1/worker/test_gpu_model_runner.py
+7
-4
@@ -680,6 +680,7 @@ def test_init_kv_cache_with_kv_sharing_valid():
680
kv_cache_spec[layer_0].page_size_bytes
681
682
runner.initialize_kv_cache(kv_cache_config)
683
+ kv_cache_config_after_init = runner.kv_cache_config
684
685
layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
686
layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
@@ -687,10 +688,12 @@ def test_init_kv_cache_with_kv_sharing_valid():
688
assert id(layer_1_kv) == id(layer_0_kv)
689
690
# check layer 1 added to kv cache group's layer names
690
- assert len(kv_cache_config.kv_cache_groups) == 1
691
- assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
692
- assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
693
- assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
691
+ assert len(kv_cache_config_after_init.kv_cache_groups) == 1
692
+ assert len(kv_cache_config_after_init.kv_cache_groups[0].layer_names) == 2
693
+ assert kv_cache_config_after_init.kv_cache_groups[0].layer_names[
694
+ 0] == layer_0
695
+ assert kv_cache_config_after_init.kv_cache_groups[0].layer_names[
696
+ 1] == layer_1
697
698
699
def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
vllm/attention/layers/chunked_local_attention.py
+16
-13
@@ -6,12 +6,13 @@ from typing import List, Optional
6
import torch
7
8
from vllm import envs
9
-from vllm.attention.backends.abstract import AttentionBackend
9
+from vllm.attention.backends.abstract import (AttentionBackend,
10
+ AttentionMetadata)
11
from vllm.attention.selector import get_attn_backend
12
from vllm.config import CacheConfig, QuantizationConfig
13
from vllm.v1.attention.backends.utils import (
14
CommonAttentionMetadata, make_local_attention_virtual_batches,
14
- subclass_attention_backend, subclass_attention_metadata_builder)
15
+ subclass_attention_backend)
16
17
from ..layer import Attention
18
@@ -24,21 +25,23 @@ def create_chunked_local_attention_backend(
25
) -> type[AttentionBackend]:
26
prefix = f"ChunkedLocalAttention_{attention_chunk_size}_{block_size}_"
27
27
- def build_preprocess_fn(cm: CommonAttentionMetadata):
28
- return make_local_attention_virtual_batches(attention_chunk_size, cm,
29
- block_size)
28
+ underlying_builder = underlying_attn_backend.get_builder_cls()
29
+
30
+ class ChunkedLocalAttentionBuilder(underlying_builder): # type: ignore
31
+
32
+ def build(self,
33
+ common_prefix_len: int,
34
+ common_attn_metadata: CommonAttentionMetadata,
35
+ fast_build: bool = False) -> AttentionMetadata:
36
+ common_attn_metadata = make_local_attention_virtual_batches(
37
+ attention_chunk_size, common_attn_metadata, block_size)
38
+ return super().build(common_prefix_len, common_attn_metadata,
39
+ fast_build)
40
31
- # Dynamically create a new attention backend that wraps the
32
- # underlying attention backend but applies
33
- # `make_local_attention_virtual_batches` before calling `build(...)`
34
- builder_cls = subclass_attention_metadata_builder(
35
- name_prefix=prefix,
36
- builder_cls=underlying_attn_backend.get_builder_cls(),
37
- build_preprocess_fn=build_preprocess_fn)
41
attn_backend = subclass_attention_backend(
42
name_prefix=prefix,
43
attention_backend_cls=underlying_attn_backend,
41
- builder_cls=builder_cls)
44
+ builder_cls=ChunkedLocalAttentionBuilder)
45
46
return attn_backend
47
vllm/attention/layers/encoder_only_attention.py
new
+86
@@ -0,0 +1,86 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
+from copy import copy
5
+from typing import Optional
6
+
7
+import torch
8
+from transformers import CacheConfig
9
+
10
+from vllm import envs
11
+from vllm.attention.backends.abstract import (AttentionBackend,
12
+ AttentionMetadata, AttentionType)
13
+from vllm.attention.layer import Attention
14
+from vllm.attention.selector import get_attn_backend
15
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
16
+ subclass_attention_backend)
17
+
18
+
19
+@functools.lru_cache
20
+def create_encoder_only_attention_backend(
21
+ underlying_attn_backend: AttentionBackend, ) -> type[AttentionBackend]:
22
+ prefix = "EncoderOnlyAttention_"
23
+ underlying_builder = underlying_attn_backend.get_builder_cls()
24
+
25
+ class EncoderOnlyAttentionBuilder(underlying_builder): # type: ignore
26
+
27
+ def build(self,
28
+ common_prefix_len: int,
29
+ common_attn_metadata: CommonAttentionMetadata,
30
+ fast_build: bool = False) -> AttentionMetadata:
31
+ new_common_attn_metadata = copy(common_attn_metadata)
32
+ new_common_attn_metadata.causal = False
33
+ return super().build(common_prefix_len, new_common_attn_metadata,
34
+ fast_build)
35
+
36
+ attn_backend = subclass_attention_backend(
37
+ name_prefix=prefix,
38
+ attention_backend_cls=underlying_attn_backend,
39
+ builder_cls=EncoderOnlyAttentionBuilder)
40
+
41
+ return attn_backend
42
+
43
+
44
+class EncoderOnlyAttention(Attention):
45
+ """
46
+ Encoder attention is a special case that doesn't need a KV Cache.
47
+ """
48
+
49
+ def __init__(self,
50
+ num_heads: int,
51
+ head_size: int,
52
+ scale: float,
53
+ cache_config: Optional[CacheConfig] = None,
54
+ attn_type: Optional[str] = None,
55
+ **kwargs):
56
+ dtype = torch.get_default_dtype()
57
+
58
+ if cache_config is not None:
59
+ kv_cache_dtype = cache_config.cache_dtype
60
+ block_size = cache_config.block_size
61
+ else:
62
+ kv_cache_dtype = "auto"
63
+ block_size = 16
64
+
65
+ if envs.VLLM_USE_V1:
66
+ underlying_attn_backend = get_attn_backend(head_size, dtype,
67
+ kv_cache_dtype,
68
+ block_size)
69
+
70
+ attn_backend = create_encoder_only_attention_backend(
71
+ underlying_attn_backend)
72
+ else:
73
+ # in v0 encoder only attention is handled inside the backends
74
+ attn_backend = None
75
+
76
+ if attn_type is not None:
77
+ assert attn_type == AttentionType.ENCODER_ONLY, \
78
+ "EncoderOnlyAttention only supports AttentionType.ENCODER_ONLY"
79
+
80
+ super().__init__(num_heads=num_heads,
81
+ head_size=head_size,
82
+ scale=scale,
83
+ cache_config=cache_config,
84
+ attn_backend=attn_backend,
85
+ attn_type=AttentionType.ENCODER_ONLY,
86
+ **kwargs)
vllm/model_executor/models/bert.py
+8
-9
@@ -8,7 +8,7 @@ import torch
8
from torch import nn
9
from transformers import BertConfig
10
11
-from vllm.attention import Attention, AttentionType
11
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, PoolerConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -239,14 +239,13 @@ class BertSelfAttention(nn.Module):
239
quant_config=quant_config,
240
prefix=f"{prefix}.qkv_proj")
241
242
- self.attn = Attention(num_heads=self.num_heads,
243
- head_size=self.head_dim,
244
- scale=self.scaling,
245
- num_kv_heads=self.num_kv_heads,
246
- cache_config=cache_config,
247
- quant_config=quant_config,
248
- prefix=f"{prefix}.attn",
249
- attn_type=AttentionType.ENCODER_ONLY)
242
+ self.attn = EncoderOnlyAttention(num_heads=self.num_heads,
243
+ head_size=self.head_dim,
244
+ scale=self.scaling,
245
+ num_kv_heads=self.num_kv_heads,
246
+ cache_config=cache_config,
247
+ quant_config=quant_config,
248
+ prefix=f"{prefix}.attn")
249
250
def forward(
251
self,
vllm/model_executor/models/bert_with_rope.py
+8
-9
@@ -7,7 +7,7 @@ import torch
7
from torch import nn
8
from transformers import PretrainedConfig
9
10
-from vllm.attention import Attention, AttentionType
10
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
@@ -119,14 +119,13 @@ class BertWithRopeAttention(nn.Module):
119
120
self.rotary_emb = get_rope(**rotary_kwargs)
121
122
- self.attn = Attention(num_heads=self.num_heads,
123
- head_size=self.head_dim,
124
- scale=self.scaling,
125
- num_kv_heads=self.num_kv_heads,
126
- cache_config=cache_config,
127
- quant_config=quant_config,
128
- prefix=f"{prefix}.attn",
129
- attn_type=AttentionType.ENCODER_ONLY)
122
+ self.attn = EncoderOnlyAttention(num_heads=self.num_heads,
123
+ head_size=self.head_dim,
124
+ scale=self.scaling,
125
+ num_kv_heads=self.num_kv_heads,
126
+ cache_config=cache_config,
127
+ quant_config=quant_config,
128
+ prefix=f"{prefix}.attn")
129
130
self.out_proj = RowParallelLinear(input_size=hidden_size,
131
output_size=hidden_size,
vllm/model_executor/models/llama.py
+5
-1
@@ -31,6 +31,7 @@ from torch import nn
31
from transformers import LlamaConfig
32
33
from vllm.attention import Attention, AttentionType
34
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -173,7 +174,10 @@ class LlamaAttention(nn.Module):
174
if is_sliding:
175
sliding_window = config.sliding_window
176
176
- self.attn = Attention(
177
+ attn_cls = (EncoderOnlyAttention
178
+ if attn_type == AttentionType.ENCODER_ONLY else Attention)
179
+
180
+ self.attn = attn_cls(
181
self.num_heads,
182
self.head_dim,
183
self.scaling,
vllm/model_executor/models/modernbert.py
+7
-7
@@ -7,7 +7,7 @@ import torch
7
from torch import nn
8
from transformers import ModernBertConfig
9
10
-from vllm.attention import Attention, AttentionType
10
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -104,12 +104,12 @@ class ModernBertAttention(nn.Module):
104
head_size=self.head_dim,
105
dim=self.head_dim,
106
base=rope_theta)
107
- self.attn = Attention(self.num_heads,
108
- self.head_dim,
109
- self.scaling,
110
- prefix=f"{layer_id}.attn",
111
- attn_type=AttentionType.ENCODER_ONLY,
112
- per_layer_sliding_window=sliding_window)
107
+ self.attn = EncoderOnlyAttention(
108
+ self.num_heads,
109
+ self.head_dim,
110
+ self.scaling,
111
+ prefix=f"{layer_id}.attn",
112
+ per_layer_sliding_window=sliding_window)
113
self.Wo = RowParallelLinear(config.hidden_size,
114
config.hidden_size,
115
bias=config.attention_bias)
vllm/model_executor/models/qwen2.py
+4
-1
@@ -32,6 +32,7 @@ from torch import nn
32
from transformers import Qwen2Config
33
34
from vllm.attention import Attention, AttentionType
35
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -159,7 +160,9 @@ class Qwen2Attention(nn.Module):
160
rope_scaling=rope_scaling,
161
dual_chunk_attention_config=dual_chunk_attention_config,
162
)
162
- self.attn = Attention(
163
+ attn_cls = (EncoderOnlyAttention
164
+ if attn_type == AttentionType.ENCODER_ONLY else Attention)
165
+ self.attn = attn_cls(
166
self.num_heads,
167
self.head_dim,
168
self.scaling,
vllm/v1/attention/backends/utils.py
+1
-31
@@ -5,8 +5,7 @@ import enum
5
import functools
6
from abc import abstractmethod
7
from dataclasses import dataclass, make_dataclass
8
-from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Generic, Optional,
9
- TypeVar)
8
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar
9
10
import numpy as np
11
import torch
@@ -543,35 +542,6 @@ def make_local_attention_virtual_batches(
542
)
543
544
546
-def subclass_attention_metadata_builder(
547
- name_prefix: str,
548
- builder_cls: type[AttentionMetadataBuilder[M]],
549
- build_preprocess_fn: Callable[[CommonAttentionMetadata],
550
- CommonAttentionMetadata],
551
-) -> type[AttentionMetadataBuilder[M]]:
552
- """
553
- Return a new subclass of `builder_cls` whose .build(...) method
554
- first calls build_preprocess_fn(common_attn_metadata) on the metadata.
555
- """
556
- name: str = name_prefix + builder_cls.__name__ # type: ignore
557
-
558
- def build(self,
559
- common_prefix_len: int,
560
- common_attn_metadata: CommonAttentionMetadata,
561
- fast_build: bool = False):
562
- return builder_cls.build(self, common_prefix_len,
563
- build_preprocess_fn(common_attn_metadata),
564
- fast_build)
565
-
566
- Wrapped = type(
567
- name,
568
- (builder_cls, ), # inherit from the original
569
- {
570
- "build": build,
571
- })
572
- return Wrapped # type: ignore
573
-
574
-
545
def subclass_attention_backend(
546
name_prefix: str, attention_backend_cls: type[AttentionBackend],
547
builder_cls: type[AttentionMetadataBuilder[M]]
vllm/v1/kv_cache_interface.py
+8
@@ -203,6 +203,14 @@ class MambaSpec(KVCacheSpec):
203
return self.page_size_bytes
204
205
206
+@dataclass(frozen=True)
207
+class EncoderOnlyAttentionSpec(AttentionSpec):
208
+
209
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
210
+ # Encoder-only layers do not need KV cache
211
+ return 0
212
+
213
+
214
@dataclass
215
class KVCacheTensor:
216
"""
vllm/v1/worker/gpu_model_runner.py
+73
-140
@@ -8,6 +8,7 @@ import time
8
from collections import defaultdict
9
from collections.abc import Iterator
10
from contextlib import contextmanager
11
+from copy import deepcopy
12
from typing import TYPE_CHECKING, Any, Optional, Union, cast
13
14
import numpy as np
@@ -62,9 +63,10 @@ from vllm.v1.attention.backends.utils import (
63
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
64
from vllm.v1.kv_cache_interface import (AttentionSpec,
65
ChunkedLocalAttentionSpec,
66
+ EncoderOnlyAttentionSpec,
67
FullAttentionSpec, KVCacheConfig,
66
- KVCacheSpec, MambaSpec,
67
- SlidingWindowSpec)
68
+ KVCacheGroupSpec, KVCacheSpec,
69
+ MambaSpec, SlidingWindowSpec)
70
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds,
71
LogprobsTensors, ModelRunnerOutput)
72
from vllm.v1.pool.metadata import PoolingMetadata
@@ -136,7 +138,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
138
cache_config.cache_dtype]
139
140
self.is_pooling_model = model_config.pooler_config is not None
139
- self.is_encoder_only_model = False
141
self.is_multimodal_raw_input_supported = (
142
model_config.is_multimodal_raw_input_supported)
143
self.max_model_len = model_config.max_model_len
@@ -345,6 +346,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
346
347
self.reorder_batch_threshold: Optional[int] = None
348
349
+ # Attention layers that are only in the KVCacheConfig of the runner
350
+ # (e.g., KV sharing, encoder-only attention), but not in the
351
+ # KVCacheConfig of the scheduler.
352
+ self.runner_only_attn_layers: set[str] = set()
353
+
354
# Cached outputs.
355
self._draft_token_ids: Optional[Union[list[list[int]],
356
torch.Tensor]] = None
@@ -834,23 +840,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
840
841
attn_metadata: dict[str, Any] = {}
842
837
- # Prepare encoder attention metadata separately
838
- # (encoder layers are not in KV cache groups)
839
- if self.is_encoder_only_model:
840
-
841
- per_layer_metadata = \
842
- self._build_encoder_only_attn_metadata(
843
- scheduler_output)
844
-
845
- # Add encoder attention metadata for all encoder layers
846
- attention_layers = get_layers_from_vllm_config(
847
- self.vllm_config, Attention)
848
- for layer_name, attn_module in attention_layers.items():
849
- if attn_module.attn_type == AttentionType.ENCODER_ONLY:
850
- common_attn_metadata, encoder_attn_metadata =\
851
- per_layer_metadata[layer_name]
852
- attn_metadata[layer_name] = encoder_attn_metadata
853
-
843
# Used in the below loop.
844
query_start_loc_cpu = self.query_start_loc_cpu[:num_reqs + 1]
845
seq_lens_cpu = self.seq_lens_cpu[:num_reqs]
@@ -863,13 +852,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
852
for kv_cache_group_id, kv_cache_group_spec in enumerate(
853
self.kv_cache_config.kv_cache_groups):
854
866
- blk_table = self.input_batch.block_table[kv_cache_group_id]
867
- blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
868
- slot_mapping = blk_table.slot_mapping[:total_num_scheduled_tokens]
869
-
870
- # Fill unused with -1. Needed for reshape_and_cache in full cuda
871
- # graph mode.
872
- blk_table.slot_mapping[total_num_scheduled_tokens:].fill_(-1)
855
+ if isinstance(kv_cache_group_spec.kv_cache_spec,
856
+ EncoderOnlyAttentionSpec):
857
+ # Encoder-only layers do not have KV cache, so we need to
858
+ # create a dummy block table and slot mapping for them.
859
+ blk_table_tensor = torch.zeros(
860
+ (num_reqs, 1),
861
+ dtype=torch.int32,
862
+ pin_memory=self.pin_memory,
863
+ device="cpu").to(self.device, non_blocking=True)
864
+ slot_mapping = torch.zeros((total_num_scheduled_tokens, ),
865
+ dtype=torch.int32,
866
+ pin_memory=self.pin_memory,
867
+ device="cpu").to(self.device,
868
+ non_blocking=True)
869
+ num_common_prefix_blocks = 0
870
+ else:
871
+ blk_table = self.input_batch.block_table[kv_cache_group_id]
872
+ blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
873
+ slot_mapping = blk_table.slot_mapping[:
874
+ total_num_scheduled_tokens]
875
+
876
+ # Fill unused with -1. Needed for reshape_and_cache in full cuda
877
+ # graph mode.
878
+ blk_table.slot_mapping[total_num_scheduled_tokens:].fill_(-1)
879
+ num_common_prefix_blocks = (
880
+ scheduler_output.
881
+ num_common_prefix_blocks[kv_cache_group_id])
882
883
common_attn_metadata = CommonAttentionMetadata(
884
query_start_loc=query_start_loc,
@@ -897,8 +906,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
906
if self.cascade_attn_enabled:
907
common_prefix_len = self._compute_cascade_attn_prefix_len(
908
num_scheduled_tokens,
900
- scheduler_output.
901
- num_common_prefix_blocks[kv_cache_group_id],
909
+ num_common_prefix_blocks,
910
kv_cache_group_spec.kv_cache_spec,
911
builder,
912
)
@@ -2812,49 +2820,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2820
# Calculate reorder batch threshold (if neeeded)
2821
self.calculate_reorder_batch_threshold()
2822
2815
- if len(self.attn_groups) > 0:
2816
- return
2817
-
2818
- # Check if model is encoder-only
2819
- block_size = self.vllm_config.cache_config.block_size
2820
- use_mla = self.vllm_config.model_config.use_mla
2821
- attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list)
2822
- for layer_name, attn_module in attn_layers.items():
2823
-
2824
- if attn_module.attn_type == AttentionType.ENCODER_ONLY:
2825
- if attn_module.sliding_window is None:
2826
- attn_spec: AttentionSpec = FullAttentionSpec(
2827
- block_size=block_size,
2828
- num_kv_heads=attn_module.num_kv_heads,
2829
- head_size=attn_module.head_size,
2830
- dtype=self.kv_cache_dtype,
2831
- use_mla=use_mla)
2832
- else:
2833
- attn_spec = SlidingWindowSpec(
2834
- block_size=block_size,
2835
- num_kv_heads=attn_module.num_kv_heads,
2836
- head_size=attn_module.head_size,
2837
- dtype=self.kv_cache_dtype,
2838
- sliding_window=attn_module.sliding_window,
2839
- use_mla=use_mla)
2840
- attn_specs[attn_spec].append(layer_name)
2841
-
2842
- else:
2843
- raise ValueError("Expected only encoder-only layers")
2844
-
2845
- if len(attn_specs) > 0:
2846
- total_layers = 0
2847
- for attn_spec, layer_names in attn_specs.items():
2848
-
2849
- attn_backends = get_attn_backends_for_layers(layer_names)
2850
- total_layers += len(layer_names)
2851
-
2852
- self.attn_groups.append(
2853
- create_attn_groups(attn_backends, attn_spec))
2854
- assert total_layers == len(attn_layers), \
2855
- "All or none of the layers are expected to be encoder-only"
2856
- self.is_encoder_only_model = True
2857
-
2823
def initialize_cudagraph_capture(self) -> None:
2824
min_cg_support = AttentionCGSupport.ALWAYS
2825
min_cg_builder_name = None
@@ -3002,7 +2967,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2967
2968
layer_names = set()
2969
for group in kv_cache_config.kv_cache_groups:
3005
- layer_names.update(group.layer_names)
2970
+ for layer_name in group.layer_names:
2971
+ if layer_name in self.runner_only_attn_layers:
2972
+ continue
2973
+ layer_names.add(layer_name)
2974
assert layer_names == set(kv_cache_raw_tensors.keys(
2975
)), "Some layers are not correctly initialized"
2976
return kv_cache_raw_tensors
@@ -3040,6 +3008,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3008
for kv_cache_spec, group in self._kv_cache_spec_attn_group_iterator():
3009
attn_backend = group.backend
3010
for layer_name in group.layer_names:
3011
+ if layer_name in self.runner_only_attn_layers:
3012
+ continue
3013
raw_tensor = kv_cache_raw_tensors[layer_name]
3014
assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
3015
num_blocks = (raw_tensor.numel() //
@@ -3161,6 +3131,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3131
kv_cache_config.kv_cache_groups,
3132
kv_caches,
3133
self.attn_groups,
3134
+ self.runner_only_attn_layers,
3135
)
3136
attn_layers = get_layers_from_vllm_config(self.vllm_config,
3137
Attention)
@@ -3185,8 +3156,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3156
kv_cache_config: Configuration for the KV cache, including the KV
3157
cache size of each layer
3158
"""
3159
+ kv_cache_config = deepcopy(kv_cache_config)
3160
self.kv_cache_config = kv_cache_config
3161
self.may_reinitialize_input_batch(kv_cache_config)
3162
+ self.may_add_encoder_only_layers_to_kv_cache_config()
3163
self.initialize_attn_backend(kv_cache_config)
3164
kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)
3165
@@ -3199,6 +3172,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3172
if has_kv_transfer_group():
3173
get_kv_transfer_group().register_kv_caches(kv_caches)
3174
3175
+ def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
3176
+ """
3177
+ Add encoder-only layers to the KV cache config.
3178
+ """
3179
+ block_size = self.vllm_config.cache_config.block_size
3180
+ use_mla = self.vllm_config.model_config.use_mla
3181
+ encoder_only_attn_specs: dict[AttentionSpec,
3182
+ list[str]] = defaultdict(list)
3183
+ attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
3184
+ for layer_name, attn_module in attn_layers.items():
3185
+ if attn_module.attn_type == AttentionType.ENCODER_ONLY:
3186
+ attn_spec = EncoderOnlyAttentionSpec(
3187
+ block_size=block_size,
3188
+ num_kv_heads=attn_module.num_kv_heads,
3189
+ head_size=attn_module.head_size,
3190
+ dtype=self.kv_cache_dtype,
3191
+ use_mla=use_mla)
3192
+ encoder_only_attn_specs[attn_spec].append(layer_name)
3193
+ self.runner_only_attn_layers.add(layer_name)
3194
+ if len(encoder_only_attn_specs) > 0:
3195
+ assert len(
3196
+ encoder_only_attn_specs
3197
+ ) == 1, "Only support one encoder-only attention spec now"
3198
+ spec, layer_names = encoder_only_attn_specs.popitem()
3199
+ self.kv_cache_config.kv_cache_groups.append(
3200
+ KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec))
3201
+
3202
def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
3203
"""
3204
Generates the KVCacheSpec by parsing the kv cache format from each
@@ -3287,70 +3287,3 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3287
mamba_type=mamba_module.mamba_type)
3288
3289
return kv_cache_spec
3290
-
3291
- def _build_encoder_only_attn_metadata(
3292
- self, scheduler_output: "SchedulerOutput") -> \
3293
- dict[str, tuple[CommonAttentionMetadata, Any]]:
3294
- """Prepare encoder attention metadata for encoder-only models.
3295
-
3296
- Args:
3297
- scheduler_output: Scheduler output
3298
-
3299
- Returns:
3300
- dict[str, Any]: Encoder attention metadata
3301
- """
3302
- num_reqs = self.input_batch.num_reqs
3303
- total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
3304
-
3305
- # Get the number of scheduled tokens for each request.
3306
- req_ids = self.input_batch.req_ids
3307
- tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
3308
- max_num_scheduled_tokens = max(tokens)
3309
-
3310
- dummy_block_table = torch.zeros((num_reqs, 1),
3311
- dtype=torch.int32,
3312
- pin_memory=self.pin_memory,
3313
- device="cpu").to(self.device,
3314
- non_blocking=True)
3315
- dummy_slot_mapping = torch.zeros((total_num_scheduled_tokens, ),
3316
- dtype=torch.int32,
3317
- pin_memory=self.pin_memory,
3318
- device="cpu").to(self.device,
3319
- non_blocking=True)
3320
-
3321
- group_metadata = dict[str, tuple[CommonAttentionMetadata, Any]]()
3322
-
3323
- for attn_group_list in self.attn_groups:
3324
-
3325
- assert len(attn_group_list) == 1
3326
- attn_group = attn_group_list[0]
3327
-
3328
- # Use the first attention metadata builder
3329
- # to create encoder attention metadata
3330
- builder = attn_group.metadata_builder
3331
-
3332
- common_metadata = CommonAttentionMetadata(
3333
- query_start_loc=self.query_start_loc[:num_reqs + 1],
3334
- query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
3335
- seq_lens=self.seq_lens[:num_reqs],
3336
- seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
3337
- num_computed_tokens_cpu=self.input_batch.
3338
- num_computed_tokens_cpu_tensor[:num_reqs],
3339
- num_reqs=num_reqs,
3340
- num_actual_tokens=total_num_scheduled_tokens,
3341
- max_query_len=max_num_scheduled_tokens,
3342
- max_seq_len=self.seq_lens_cpu[:num_reqs].max().item(),
3343
- block_table_tensor=dummy_block_table,
3344
- slot_mapping=dummy_slot_mapping,
3345
- causal=False,
3346
- )
3347
-
3348
- metadata = builder.build(
3349
- common_prefix_len=0, # No cascade for encoder
3350
- common_attn_metadata=common_metadata,
3351
- )
3352
-
3353
- for layer_name in attn_group.layer_names:
3354
- group_metadata[layer_name] = (common_metadata, metadata)
3355
-
3356
- return group_metadata
vllm/v1/worker/utils.py
+4
@@ -204,6 +204,7 @@ def initialize_kv_cache_for_kv_sharing(
204
kv_caches: dict[str, torch.Tensor],
205
# Optional for now to avoid breaking TPU
206
attn_groups: Optional[list[list[AttentionGroup]]] = None,
207
+ runner_only_attn_layers: Optional[set[str]] = None,
208
) -> None:
209
"""
210
Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
@@ -250,6 +251,9 @@ def initialize_kv_cache_for_kv_sharing(
251
attn_groups[kv_cache_group_idx][attn_group_idx].layer_names.append(
252
layer_name)
253
254
+ if runner_only_attn_layers is not None:
255
+ runner_only_attn_layers.add(layer_name)
256
+
257
258
def bind_kv_cache(
259
kv_caches: dict[str, torch.Tensor],