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] Support multiple attention metadata builders per kv_cache_spec + proper local attention no hybrid kv cache fix (#21588)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Aug 6, 2025 at 21:40 UTC
1dc8a70b6d4e8ba4e139f1ddb86a166694f42f21
13 files changed
+368
-212
tests/v1/spec_decode/test_eagle.py
+2
-1
@@ -313,7 +313,8 @@ def test_propose(num_speculative_tokens, backend):
313
314
# Mock runner for attention metadata building
315
proposer.runner = mock.MagicMock()
316
- proposer.runner.attn_metadata_builders = [attn_metadata_builder]
316
+ proposer.runner.attn_groups.append([mock.MagicMock()])
317
+ proposer.runner.attn_groups[0][0].metadata_builder = attn_metadata_builder
318
319
result = proposer.propose(target_token_ids=target_token_ids,
320
target_positions=target_positions,
tests/v1/worker/test_gpu_model_runner.py
+3
-3
@@ -417,12 +417,12 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
417
return rnd_stride
418
419
# Patch the attention backend class and re-trigger the KV cache creation.
420
- for attn_backend in model_runner.attn_backends:
420
+ for attn_group in model_runner._attn_group_iterator():
421
+ attn_backend = attn_group.backend
422
monkeypatch.setattr(attn_backend, "get_kv_cache_stride_order",
423
rnd_stride_order)
424
424
- model_runner.attn_backends = []
425
- model_runner.attn_metadata_builders = []
425
+ model_runner.attn_groups = []
426
model_runner.initialize_kv_cache(model_runner.kv_cache_config)
427
428
# Shape is unchanged, but layout may differ
vllm/attention/backends/abstract.py
+4
@@ -106,6 +106,10 @@ class AttentionBackend(ABC):
106
block_size: int, num_seqs: int, num_queries: int) -> None:
107
raise NotImplementedError
108
109
+ @classmethod
110
+ def full_cls_name(cls) -> tuple[str, str]:
111
+ return (cls.__module__, cls.__qualname__)
112
+
113
114
@dataclass
115
class AttentionMetadata:
vllm/attention/layer.py
+18
-18
@@ -9,6 +9,7 @@ import torch.nn.functional as F
9
10
import vllm.envs as envs
11
from vllm.attention import AttentionType
12
+from vllm.attention.backends.abstract import AttentionBackend
13
from vllm.attention.selector import backend_name_to_enum, get_attn_backend
14
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
15
from vllm.config import CacheConfig, get_current_vllm_config
@@ -80,6 +81,7 @@ class Attention(nn.Module):
81
prefix: str = "",
82
attn_type: str = AttentionType.DECODER,
83
kv_sharing_target_layer_name: Optional[str] = None,
84
+ attn_backend: Optional[type[AttentionBackend]] = None,
85
**extra_impl_args,
86
) -> None:
87
"""
@@ -137,15 +139,6 @@ class Attention(nn.Module):
139
self.num_kv_heads = num_kv_heads
140
self.sliding_window = sliding_window
141
140
- # For v1 we have backend agnostic iRoPE (local chunked attention)
141
- # we have to store the flag on the layer so gpu model runner can
142
- # set KVSpec appropriately (and pop it so it doesnt get passed to
143
- # the backends)
144
- if envs.VLLM_USE_V1:
145
- self.use_irope = extra_impl_args.pop("use_irope", False)
146
- else:
147
- self.use_irope = extra_impl_args.get("use_irope", False)
148
-
142
quant_method = quant_config.get_quant_method(
143
self, prefix=prefix) if quant_config else None
144
if quant_method is not None and not isinstance(
@@ -166,18 +159,22 @@ class Attention(nn.Module):
159
# During model initialization, the default dtype is set as the model
160
# weight and activation dtype.
161
dtype = torch.get_default_dtype()
169
- attn_backend = get_attn_backend(head_size,
170
- dtype,
171
- kv_cache_dtype,
172
- block_size,
173
- is_attention_free,
174
- use_mla=use_mla)
175
- impl_cls = attn_backend.get_impl_cls()
162
+ if attn_backend is None:
163
+ self.attn_backend = get_attn_backend(head_size,
164
+ dtype,
165
+ kv_cache_dtype,
166
+ block_size,
167
+ is_attention_free,
168
+ use_mla=use_mla)
169
+ else:
170
+ self.attn_backend = attn_backend
171
+
172
+ impl_cls = self.attn_backend.get_impl_cls()
173
self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
174
alibi_slopes, sliding_window, kv_cache_dtype,
175
logits_soft_cap, attn_type,
176
kv_sharing_target_layer_name, **extra_impl_args)
180
- self.backend = backend_name_to_enum(attn_backend.get_name())
177
+ self.backend = backend_name_to_enum(self.attn_backend.get_name())
178
self.dtype = dtype
179
180
# For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
@@ -187,7 +184,7 @@ class Attention(nn.Module):
184
self.use_direct_call = not current_platform.is_cuda_alike(
185
) and not current_platform.is_cpu()
186
190
- self.use_output = attn_backend.accept_output_buffer
187
+ self.use_output = self.attn_backend.accept_output_buffer
188
compilation_config = get_current_vllm_config().compilation_config
189
if prefix in compilation_config.static_forward_context:
190
raise ValueError(f"Duplicate layer name: {prefix}")
@@ -309,6 +306,9 @@ class Attention(nn.Module):
306
if hasattr(self.impl, "process_weights_after_loading"):
307
self.impl.process_weights_after_loading(act_dtype)
308
309
+ def get_attn_backend(self) -> type[AttentionBackend]:
310
+ return self.attn_backend
311
+
312
313
class MultiHeadAttention(nn.Module):
314
"""Multi-headed attention without any cache, used for ViT."""
vllm/attention/layers/chunked_local_attention.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
+from typing import List, Optional
5
+
6
+import torch
7
+
8
+from vllm import envs
9
+from vllm.attention.backends.abstract import AttentionBackend
10
+from vllm.attention.selector import get_attn_backend
11
+from vllm.config import CacheConfig, QuantizationConfig
12
+from vllm.v1.attention.backends.utils import (
13
+ CommonAttentionMetadata, make_local_attention_virtual_batches,
14
+ subclass_attention_backend, subclass_attention_metadata_builder)
15
+
16
+from ..layer import Attention
17
+
18
+
19
+@functools.lru_cache
20
+def create_chunked_local_attention_backend(
21
+ underlying_attn_backend: AttentionBackend,
22
+ attention_chunk_size: int,
23
+ block_size: int,
24
+) -> type[AttentionBackend]:
25
+ prefix = f"ChunkedLocalAttention_{attention_chunk_size}_{block_size}_"
26
+
27
+ def build_preprocess_fn(cm: CommonAttentionMetadata):
28
+ return make_local_attention_virtual_batches(attention_chunk_size, cm,
29
+ block_size)
30
+
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)
38
+ attn_backend = subclass_attention_backend(
39
+ name_prefix=prefix,
40
+ attention_backend_cls=underlying_attn_backend,
41
+ builder_cls=builder_cls)
42
+
43
+ return attn_backend
44
+
45
+
46
+class ChunkedLocalAttention(Attention):
47
+
48
+ def __init__(self,
49
+ num_heads: int,
50
+ head_size: int,
51
+ scale: float,
52
+ attention_chunk_size: int,
53
+ num_kv_heads: Optional[int] = None,
54
+ alibi_slopes: Optional[List[float]] = None,
55
+ cache_config: Optional[CacheConfig] = None,
56
+ quant_config: Optional[QuantizationConfig] = None,
57
+ kv_sharing_target_layer_name: Optional[str] = None,
58
+ prefix: str = ""):
59
+ dtype = torch.get_default_dtype()
60
+ if cache_config is not None:
61
+ kv_cache_dtype = cache_config.cache_dtype
62
+ block_size = cache_config.block_size
63
+ else:
64
+ kv_cache_dtype = "auto"
65
+ block_size = 16
66
+
67
+ if envs.VLLM_USE_V1:
68
+ underlying_attn_backend = get_attn_backend(head_size, dtype,
69
+ kv_cache_dtype,
70
+ block_size)
71
+
72
+ attn_backend = create_chunked_local_attention_backend(
73
+ underlying_attn_backend, attention_chunk_size, block_size)
74
+ else:
75
+ # in v0 the local attention is handled inside the backends
76
+ attn_backend = None
77
+
78
+ super().__init__(
79
+ num_heads=num_heads,
80
+ head_size=head_size,
81
+ scale=scale,
82
+ num_kv_heads=num_kv_heads,
83
+ alibi_slopes=alibi_slopes,
84
+ cache_config=cache_config,
85
+ quant_config=quant_config,
86
+ prefix=prefix,
87
+ kv_sharing_target_layer_name=kv_sharing_target_layer_name,
88
+ attn_backend=attn_backend)
vllm/attention/selector.py
+1
-1
@@ -142,7 +142,7 @@ def get_attn_backend(
142
dtype: torch.dtype,
143
kv_cache_dtype: Optional[str],
144
block_size: int,
145
- is_attention_free: bool,
145
+ is_attention_free: bool = False,
146
use_mla: bool = False,
147
) -> type[AttentionBackend]:
148
"""Selects which attention backend to use and lazily imports it."""
vllm/model_executor/models/llama4.py
+6
-4
@@ -25,6 +25,7 @@ from torch import nn
25
from transformers import Llama4TextConfig
26
27
from vllm.attention import Attention
28
+from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
29
from vllm.compilation.decorators import support_torch_compile
30
from vllm.config import CacheConfig, VllmConfig
31
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -194,17 +195,18 @@ class Llama4Attention(nn.Module):
195
is_neox_style=is_neox_style,
196
) if not self.nope else None
197
197
- self.attn = Attention(
198
+ attn_cls = Attention if self.nope else ChunkedLocalAttention
199
+ self.attn = attn_cls(
200
self.num_heads,
201
self.head_dim,
202
self.scaling,
203
num_kv_heads=self.num_kv_heads,
204
cache_config=cache_config,
205
quant_config=quant_config,
204
- per_layer_sliding_window=None,
205
- use_irope=not self.nope,
206
prefix=f"{prefix}.attn",
207
- )
207
+ **({
208
+ "attention_chunk_size": config.attention_chunk_size
209
+ } if not self.nope else {}))
210
211
def _get_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
212
floor = torch.floor((positions + 1.0) / self.floor_scale)
vllm/v1/attention/backends/utils.py
+46
-2
@@ -5,12 +5,12 @@ import enum
5
import functools
6
from abc import abstractmethod
7
from dataclasses import dataclass, make_dataclass
8
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar
8
+from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Generic, Optional,
9
+ TypeVar)
10
11
import numpy as np
12
import torch
13
13
-from vllm.attention.layer import Attention
14
from vllm.config import VllmConfig, get_layers_from_vllm_config
15
from vllm.utils import cdiv
16
@@ -20,6 +20,8 @@ if TYPE_CHECKING:
20
from vllm.v1.worker.gpu_input_batch import InputBatch
21
22
import vllm.envs as envs
23
+from vllm.attention.backends.abstract import AttentionBackend
24
+from vllm.attention.layer import Attention
25
from vllm.distributed.kv_transfer.kv_connector.utils import (
26
get_kv_connector_cache_layout)
27
from vllm.logger import init_logger
@@ -532,6 +534,48 @@ def make_local_attention_virtual_batches(
534
)
535
536
537
+def subclass_attention_metadata_builder(
538
+ name_prefix: str,
539
+ builder_cls: type[AttentionMetadataBuilder[M]],
540
+ build_preprocess_fn: Callable[[CommonAttentionMetadata],
541
+ CommonAttentionMetadata],
542
+) -> type[AttentionMetadataBuilder[M]]:
543
+ """
544
+ Return a new subclass of `builder_cls` whose .build(...) method
545
+ first calls build_preprocess_fn(common_attn_metadata) on the metadata.
546
+ """
547
+ name: str = name_prefix + builder_cls.__name__ # type: ignore
548
+
549
+ def build(self,
550
+ common_prefix_len: int,
551
+ common_attn_metadata: CommonAttentionMetadata,
552
+ fast_build: bool = False):
553
+ return builder_cls.build(self, common_prefix_len,
554
+ build_preprocess_fn(common_attn_metadata),
555
+ fast_build)
556
+
557
+ Wrapped = type(
558
+ name,
559
+ (builder_cls, ), # inherit from the original
560
+ {
561
+ "build": build,
562
+ })
563
+ return Wrapped # type: ignore
564
+
565
+
566
+def subclass_attention_backend(
567
+ name_prefix: str, attention_backend_cls: type[AttentionBackend],
568
+ builder_cls: type[AttentionMetadataBuilder[M]]
569
+) -> type[AttentionBackend]:
570
+ """
571
+ Return a new subclass where `get_builder_cls` returns `builder_cls`.
572
+ """
573
+ name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
574
+
575
+ return type(name, (attention_backend_cls, ),
576
+ {"get_builder_cls": lambda: builder_cls})
577
+
578
+
579
def split_decodes_and_prefills(
580
common_attn_metadata: CommonAttentionMetadata,
581
decode_threshold: int = 1,
vllm/v1/spec_decode/eagle.py
+5
-4
@@ -158,9 +158,9 @@ class EagleProposer:
158
assert self.runner is not None
159
160
# FIXME: need to consider multiple kv_cache_groups
161
- attn_metadata = self.runner.attn_metadata_builders[
162
- 0].build_for_drafting(common_attn_metadata=common_attn_metadata,
163
- draft_index=0)
161
+ attn_metadata = self.runner.attn_groups[0][0].metadata_builder\
162
+ .build_for_drafting(common_attn_metadata=common_attn_metadata,
163
+ draft_index=0)
164
165
# At this moment, we assume all eagle layers belong to the same KV
166
# cache group, thus using the same attention metadata.
@@ -349,7 +349,8 @@ class EagleProposer:
349
hidden_states: torch.Tensor,
350
common_attn_metadata: CommonAttentionMetadata,
351
) -> list[torch.Tensor]:
352
- tree_attn_metadata_builder = self.runner.attn_metadata_builders[0]
352
+ tree_attn_metadata_builder = \
353
+ self.runner.attn_groups[0][0].metadata_builder
354
assert isinstance(tree_attn_metadata_builder,
355
TreeAttentionMetadataBuilder)
356
vllm/v1/worker/cpu_model_runner.py
+4
-4
@@ -53,11 +53,11 @@ class CPUModelRunner(GPUModelRunner):
53
raise ValueError("Multiple KVCacheGroups is not"
54
"currently supported with CPU model runner.")
55
56
- assert type(
57
- self.attn_metadata_builders[0]) is TorchSDPAMetadataBuilderV1
56
+ assert type(self.attn_groups[0]
57
+ [0].metadata_builder) is TorchSDPAMetadataBuilderV1
58
59
- self.attn_metadata_builders[0].reorder_batch(self.input_batch,
60
- scheduler_output)
59
+ self.attn_groups[0][0].metadata_builder.reorder_batch(
60
+ self.input_batch, scheduler_output)
61
62
def _postprocess_tenosrs(self) -> None:
63
# Note: replace device tensors with cpu tensors
vllm/v1/worker/gpu_model_runner.py
+167
-173
@@ -3,7 +3,10 @@
3
4
import dataclasses
5
import gc
6
+import itertools
7
import time
8
+from collections import defaultdict
9
+from collections.abc import Iterator
10
from contextlib import contextmanager
11
from typing import TYPE_CHECKING, Any, Optional, Union, cast
12
@@ -14,9 +17,9 @@ import torch.nn as nn
17
from tqdm import tqdm
18
19
import vllm.envs as envs
17
-from vllm.attention import AttentionType, get_attn_backend
20
+from vllm.attention import Attention, AttentionType
21
from vllm.attention.backends.abstract import AttentionBackend
19
-from vllm.attention.layer import Attention
22
+from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
23
from vllm.compilation.counter import compilation_counter
24
from vllm.config import (CompilationLevel, VllmConfig,
25
get_layers_from_vllm_config, update_config)
@@ -50,7 +53,6 @@ from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
53
from vllm.v1.attention.backends.utils import (
54
AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
55
make_kv_sharing_fast_prefill_attention_metadata,
53
- make_local_attention_virtual_batches,
56
reorder_batch_to_split_decodes_and_prefills)
57
from vllm.v1.kv_cache_interface import (AttentionSpec,
58
ChunkedLocalAttentionSpec,
@@ -73,8 +75,8 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (
75
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
76
77
from ..sample.logits_processor import LogitsProcessorManager
76
-from .utils import (MultiModalBudget, bind_kv_cache, gather_mm_placeholders,
77
- initialize_kv_cache_for_kv_sharing,
78
+from .utils import (AttentionGroup, MultiModalBudget, bind_kv_cache,
79
+ gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
80
sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
81
82
if TYPE_CHECKING:
@@ -162,8 +164,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
164
# self.model: nn.Module # Set after load_model
165
# Initialize in initialize_kv_cache
166
self.kv_caches: list[torch.Tensor] = []
165
- self.attn_metadata_builders: list[AttentionMetadataBuilder] = []
166
- self.attn_backends: list[type[AttentionBackend]] = []
167
+ # indexes: [kv_cache_group_id][attn_group]
168
+ self.attn_groups: list[list[AttentionGroup]] = []
169
# self.kv_cache_config: KVCacheConfig
170
171
# req_id -> (input_id -> encoder_output)
@@ -830,81 +832,51 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
832
spec_decode_common_attn_metadata is None:
833
spec_decode_common_attn_metadata = common_attn_metadata
834
833
- if isinstance(kv_cache_group_spec.kv_cache_spec,
834
- ChunkedLocalAttentionSpec):
835
- common_attn_metadata = make_local_attention_virtual_batches(
836
- kv_cache_group_spec.kv_cache_spec.attention_chunk_size,
837
- common_attn_metadata, self.cache_config.block_size)
838
-
839
- # Prepare for cascade attention if enabled & beneficial.
840
- common_prefix_len = 0
841
- builder = self.attn_metadata_builders[kv_cache_group_id]
842
- if self.cascade_attn_enabled:
843
- common_prefix_len = self._compute_cascade_attn_prefix_len(
844
- num_scheduled_tokens,
845
- scheduler_output.
846
- num_common_prefix_blocks[kv_cache_group_id],
847
- kv_cache_group_spec.kv_cache_spec,
848
- builder,
849
- )
850
-
851
- attn_metadata_i = (builder.build(
852
- common_prefix_len=common_prefix_len,
853
- common_attn_metadata=common_attn_metadata,
854
- ))
855
-
856
- fast_prefill_metadata = attn_metadata_i
857
- if (self.cache_config.kv_sharing_fast_prefill
858
- and self.kv_sharing_fast_prefill_eligible_layers):
859
- # Dynamically create a a dataclass type that inherits
860
- # from attention metadata type but includes additional
861
- # fields logits_indices_padded and num_logits_indices
862
- # which are required for prefill truncation
863
- fast_prefill_metadata_type = (
864
- make_kv_sharing_fast_prefill_attention_metadata(
865
- metadata_cls=type(attn_metadata_i), ))
866
- fast_prefill_metadata = fast_prefill_metadata_type(
867
- **dataclasses.asdict(attn_metadata_i),
868
- logits_indices_padded=logits_indices_padded,
869
- num_logits_indices=logits_indices.size(0),
870
- )
871
-
872
- for layer_name in kv_cache_group_spec.layer_names:
873
- if (self.cache_config.kv_sharing_fast_prefill and layer_name
874
- in self.kv_sharing_fast_prefill_eligible_layers):
875
- attn_metadata[layer_name] = fast_prefill_metadata
876
- continue
835
+ for attn_group in self.attn_groups[kv_cache_group_id]:
836
+ # Prepare for cascade attention if enabled & beneficial.
837
+ common_prefix_len = 0
838
+ builder = attn_group.metadata_builder
839
+ if self.cascade_attn_enabled:
840
+ common_prefix_len = self._compute_cascade_attn_prefix_len(
841
+ num_scheduled_tokens,
842
+ scheduler_output.
843
+ num_common_prefix_blocks[kv_cache_group_id],
844
+ kv_cache_group_spec.kv_cache_spec,
845
+ builder,
846
+ )
847
878
- attn_metadata[layer_name] = attn_metadata_i
879
-
880
- # Hack for now to fix chunked local attention + no hybrid kv cache
881
- # manager we can remove this once
882
- # https://github.com/vllm-project/vllm/pull/21588
883
- # is merged (i.e. properly handle different attention backends for
884
- # the same kv_cache_spec)
885
- if self.attention_chunk_size is not None \
886
- and self.scheduler_config.disable_hybrid_kv_cache_manager:
887
- if not hasattr(self, "local_attention_layers"):
888
- self.local_attention_layers = []
889
- attn_layers = get_layers_from_vllm_config(
890
- self.vllm_config, Attention)
891
- for layer_name, attn_module in attn_layers.items():
892
- if attn_module.use_irope:
893
- self.local_attention_layers.append(layer_name)
894
-
895
- local_attn_metadata_i = (builder.build(
896
- common_prefix_len=0,
897
- common_attn_metadata=make_local_attention_virtual_batches(
898
- self.attention_chunk_size, common_attn_metadata,
899
- self.cache_config.block_size),
848
+ attn_metadata_i = (builder.build(
849
+ common_prefix_len=common_prefix_len,
850
+ common_attn_metadata=common_attn_metadata,
851
))
852
902
- for layer_name in self.local_attention_layers:
903
- attn_metadata[layer_name] = local_attn_metadata_i
853
+ fast_prefill_metadata = attn_metadata_i
854
+ if (self.cache_config.kv_sharing_fast_prefill
855
+ and self.kv_sharing_fast_prefill_eligible_layers):
856
+ # Dynamically create a a dataclass type that inherits
857
+ # from attention metadata type but includes additional
858
+ # fields logits_indices_padded and num_logits_indices
859
+ # which are required for prefill truncation
860
+ fast_prefill_metadata_type = (
861
+ make_kv_sharing_fast_prefill_attention_metadata(
862
+ metadata_cls=type(attn_metadata_i), ))
863
+ fast_prefill_metadata = fast_prefill_metadata_type(
864
+ **dataclasses.asdict(attn_metadata_i),
865
+ logits_indices_padded=logits_indices_padded,
866
+ num_logits_indices=logits_indices.size(0),
867
+ )
868
+
869
+ for layer_name in attn_group.layer_names:
870
+ if (self.cache_config.kv_sharing_fast_prefill
871
+ and layer_name
872
+ in self.kv_sharing_fast_prefill_eligible_layers):
873
+ attn_metadata[layer_name] = fast_prefill_metadata
874
+ continue
875
+ attn_metadata[layer_name] = attn_metadata_i
876
877
attention_cuda_graphs = all(
906
- b.can_run_in_cudagraph(common_attn_metadata)
907
- for b in self.attn_metadata_builders)
878
+ g.metadata_builder.can_run_in_cudagraph(common_attn_metadata)
879
+ for g in self._attn_group_iterator())
880
881
# Hot-Swap lora model
882
if self.lora_config:
@@ -2229,11 +2201,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2201
block_table[kv_cache_group_id].slot_mapping[:num_tokens],
2202
causal=True)
2203
2232
- attn_metadata_i = self.attn_metadata_builders[
2233
- kv_cache_group_id].build_for_cudagraph_capture(
2234
- common_attn_metadata)
2235
- for layer_name in kv_cache_group_spec.layer_names:
2236
- attn_metadata[layer_name] = attn_metadata_i
2204
+ for attn_group in self.attn_groups[kv_cache_group_id]:
2205
+ attn_metadata_i = attn_group.metadata_builder\
2206
+ .build_for_cudagraph_capture(common_attn_metadata)
2207
+ for layer_name in kv_cache_group_spec.layer_names:
2208
+ attn_metadata[layer_name] = attn_metadata_i
2209
2210
with self.maybe_dummy_run_with_lora(self.lora_config,
2211
num_scheduled_tokens):
@@ -2565,88 +2537,100 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2537
logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
2538
elapsed_time, cuda_graph_size / (1 << 30))
2539
2568
- def _initialize_single_attn_backend(
2569
- self, kv_cache_spec: KVCacheSpec, layer_names: list[str]
2570
- ) -> tuple[AttentionBackend, AttentionMetadataBuilder]:
2571
- if isinstance(kv_cache_spec, AttentionSpec):
2572
- attn_backend_i = get_attn_backend(
2573
- kv_cache_spec.head_size,
2574
- self.dtype,
2575
- kv_cache_spec.dtype,
2576
- kv_cache_spec.block_size,
2577
- self.model_config.is_attention_free,
2578
- use_mla=kv_cache_spec.use_mla,
2579
- )
2580
- if attn_backend_i is None:
2581
- error_msg = (f"Error with get_attn_backend: "
2582
- f"{kv_cache_spec.head_size=}, "
2583
- f"{self.dtype=}, {kv_cache_spec.dtype=}, "
2584
- f"{kv_cache_spec.block_size=}, "
2585
- f"{self.model_config.is_attention_free=}, "
2586
- f"{kv_cache_spec.use_mla=}")
2587
- logger.error(error_msg)
2588
- raise NotImplementedError(
2589
- "Non-Attention backend is not supported by V1 "
2590
- "GPUModelRunner.")
2591
- elif isinstance(kv_cache_spec, MambaSpec):
2592
- attn_backend_i = get_mamba_attn_backend(kv_cache_spec.mamba_type)
2593
- else:
2594
- raise ValueError(
2595
- f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2596
-
2597
- attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2598
- kv_cache_spec,
2599
- layer_names,
2600
- self.vllm_config,
2601
- self.device,
2602
- )
2603
-
2604
- if self.full_cuda_graph:
2605
- if attn_metadata_builder_i.attn_cudagraph_support == \
2606
- AttentionCGSupport.NEVER:
2607
- raise ValueError(f"Full CUDAGraph not supported for "
2608
- f"{attn_backend_i.__name__}. Turn off "
2609
- f"CompilationConfig.full_cuda_graph or use a "
2610
- f" different attention backend.")
2611
- if attn_metadata_builder_i.attn_cudagraph_support == \
2612
- AttentionCGSupport.PURE_DECODE_ONLY:
2613
- # Limit the max cudagraph size to the max number of
2614
- # sequences for pure decode only cudagraph backend,
2615
- # whose max_query_len is 1.
2616
- self.cudagraph_batch_sizes = [
2617
- size for size in self.cudagraph_batch_sizes
2618
- if size <= self.scheduler_config.max_num_seqs
2619
- ]
2620
- return attn_backend_i, attn_metadata_builder_i
2621
-
2540
def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
2541
"""
2542
Initialize the attention backends and attention metadata builders.
2543
"""
2626
- assert len(self.attn_backends) == 0 and len(
2627
- self.attn_metadata_builders
2628
- ) == 0, "Attention backends are already initialized"
2629
- for i, kv_cache_group_spec in enumerate(
2630
- kv_cache_config.kv_cache_groups):
2544
+ assert len(self.attn_groups) == 0, \
2545
+ "Attention backends are already initialized"
2546
+ attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
2547
+
2548
+ def get_attn_backends_for_layers(
2549
+ layer_names: list[str]
2550
+ ) -> dict[type[AttentionBackend], list[str]]:
2551
+ attn_backends = {}
2552
+ attn_backend_layers = defaultdict(list)
2553
+ # Dedupe based on full class name; this is a bit safer than using
2554
+ # using the class itself as the key because when we create dynamic
2555
+ # attention backend subclasses (e.g. ChunkedLocalAttention) unless
2556
+ # they are cached correctly, there will be different objects per
2557
+ # layer.
2558
+ for layer_name in layer_names:
2559
+ attn_backend = attn_layers[layer_name].get_attn_backend()
2560
+ key = attn_backend.full_cls_name()
2561
+ attn_backends[key] = attn_backend
2562
+ attn_backend_layers[key].append(layer_name)
2563
+ return {
2564
+ attn_backends[k]: v
2565
+ for k, v in attn_backend_layers.items()
2566
+ }
2567
+
2568
+ def create_attn_groups(
2569
+ attn_backends_map: dict[AttentionBackend, list[str]],
2570
+ kv_cache_spec: KVCacheSpec,
2571
+ ) -> list[AttentionGroup]:
2572
+ attn_groups: list[AttentionGroup] = []
2573
+ for attn_backend, layer_names in attn_backends_map.items():
2574
+ attn_metadata_builder_i = attn_backend.get_builder_cls()(
2575
+ kv_cache_spec,
2576
+ layer_names,
2577
+ self.vllm_config,
2578
+ self.device,
2579
+ )
2580
+ attn_group = AttentionGroup(attn_backend,
2581
+ attn_metadata_builder_i,
2582
+ layer_names)
2583
+ attn_groups.append(attn_group)
2584
+
2585
+ if self.full_cuda_graph:
2586
+ if attn_metadata_builder_i.attn_cudagraph_support == \
2587
+ AttentionCGSupport.NEVER:
2588
+ raise ValueError(
2589
+ f"Full CUDAGraph not supported for "
2590
+ f"{attn_backend.__name__}. Turn off "
2591
+ f"CompilationConfig.full_cuda_graph or use a "
2592
+ f" different attention backend.")
2593
+ if attn_metadata_builder_i.attn_cudagraph_support == \
2594
+ AttentionCGSupport.PURE_DECODE_ONLY:
2595
+ # Limit the max cudagraph size to the max number of
2596
+ # sequences for pure decode only cudagraph backend,
2597
+ # whose max_query_len is 1.
2598
+ self.cudagraph_batch_sizes = [
2599
+ size for size in self.cudagraph_batch_sizes
2600
+ if size <= self.scheduler_config.max_num_seqs
2601
+ ]
2602
+
2603
+ return attn_groups
2604
+
2605
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
2606
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2607
+ if isinstance(kv_cache_spec, AttentionSpec):
2608
+ attn_backends = get_attn_backends_for_layers(
2609
+ kv_cache_group_spec.layer_names)
2610
+ # TODO(lucas): move `get_mamba_attn_backend` into the mamba
2611
+ # layers like above
2612
+ elif isinstance(kv_cache_spec, MambaSpec):
2613
+ attn_backends = {
2614
+ get_mamba_attn_backend(kv_cache_spec.mamba_type):
2615
+ kv_cache_group_spec.layer_names
2616
+ }
2617
+ else:
2618
+ raise ValueError(
2619
+ f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2620
2633
- attn_backend_i, attn_metadata_builder_i = (
2634
- self._initialize_single_attn_backend(
2635
- kv_cache_spec, kv_cache_group_spec.layer_names))
2636
- self.attn_backends.append(attn_backend_i)
2637
- self.attn_metadata_builders.append(attn_metadata_builder_i)
2621
+ self.attn_groups.append(
2622
+ create_attn_groups(attn_backends, kv_cache_spec))
2623
2624
# Calculate reorder batch threshold (if neeeded)
2625
self.calculate_reorder_batch_threshold()
2626
2642
- if len(self.attn_backends) > 0:
2627
+ if len(self.attn_groups) > 0:
2628
return
2629
2630
# Check if model is encoder-only
2631
block_size = self.vllm_config.cache_config.block_size
2632
use_mla = self.vllm_config.model_config.use_mla
2633
attn_specs = list[AttentionSpec]()
2649
- attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
2634
for attn_module in attn_layers.values():
2635
2636
if attn_module.attn_type == AttentionType.ENCODER_ONLY:
@@ -2666,11 +2650,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2650
assert len(attn_specs) == len(attn_layers), \
2651
"All or none of the layers are expected to be encoder-only"
2652
2669
- attn_backend, attn_metadata_builder = (
2670
- self._initialize_single_attn_backend(attn_specs[0],
2671
- attn_layers.keys()))
2672
- self.attn_backends.append(attn_backend)
2673
- self.attn_metadata_builders.append(attn_metadata_builder)
2653
+ attn_backends = get_attn_backends_for_layers(attn_layers.keys())
2654
+
2655
+ self.attn_groups.append(
2656
+ create_attn_groups(attn_backends, attn_specs[0]))
2657
self.is_encoder_only_model = True
2658
2659
def calculate_reorder_batch_threshold(self) -> None:
@@ -2678,7 +2661,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2661
Check that if any backends reorder batches; that the reordering
2662
is compatible (e.g., decode threshold is the same)
2663
"""
2681
- for attn_metadata_builder_i in self.attn_metadata_builders:
2664
+ for group in self._attn_group_iterator():
2665
+ attn_metadata_builder_i = group.metadata_builder
2666
+
2667
# check that if any backends reorder batches; that the reordering
2668
# is compatible (e.g., decode threshold is the same)
2669
reorder_batch_threshold_i = (
@@ -2752,6 +2737,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2737
)), "Some layers are not correctly initialized"
2738
return kv_cache_raw_tensors
2739
2740
+ def _attn_group_iterator(self) -> Iterator[AttentionGroup]:
2741
+ return itertools.chain.from_iterable(self.attn_groups)
2742
+
2743
+ def _kv_cache_spec_attn_group_iterator(
2744
+ self) -> Iterator[tuple[KVCacheSpec, AttentionGroup]]:
2745
+ if not self.kv_cache_config.kv_cache_groups:
2746
+ return
2747
+ for kv_cache_spec_id, attn_groups in enumerate(self.attn_groups):
2748
+ for attn_group in attn_groups:
2749
+ yield self.kv_cache_config.kv_cache_groups[
2750
+ kv_cache_spec_id].kv_cache_spec, attn_group
2751
+
2752
def _reshape_kv_cache_tensors(
2753
self,
2754
kv_cache_config: KVCacheConfig,
@@ -2770,23 +2767,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2767
"""
2768
kv_caches: dict[str, torch.Tensor] = {}
2769
has_attn, has_mamba = False, False
2773
- for i, kv_cache_group_spec in enumerate(
2774
- kv_cache_config.kv_cache_groups):
2775
- kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2776
- for layer_name in kv_cache_group_spec.layer_names:
2770
+ for kv_cache_spec, group in self._kv_cache_spec_attn_group_iterator():
2771
+ attn_backend = group.backend
2772
+ for layer_name in group.layer_names:
2773
raw_tensor = kv_cache_raw_tensors[layer_name]
2774
assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
2775
num_blocks = (raw_tensor.numel() //
2776
kv_cache_spec.page_size_bytes)
2777
if isinstance(kv_cache_spec, AttentionSpec):
2778
has_attn = True
2783
- kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
2779
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
2780
num_blocks, kv_cache_spec.block_size,
2781
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
2782
dtype = kv_cache_spec.dtype
2783
try:
2788
- kv_cache_stride_order = self.attn_backends[
2789
- i].get_kv_cache_stride_order()
2784
+ kv_cache_stride_order = \
2785
+ attn_backend.get_kv_cache_stride_order()
2786
assert len(kv_cache_stride_order) == len(
2787
kv_cache_shape)
2788
except (AttributeError, NotImplementedError):
@@ -2850,15 +2846,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2846
kv_cache_raw_tensors: The KV cache buffer of each layer.
2847
"""
2848
2853
- for i, kv_cache_group_spec in enumerate(
2854
- kv_cache_config.kv_cache_groups):
2855
- kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2856
- for layer_name in kv_cache_group_spec.layer_names:
2849
+ for kv_cache_spec, group in self._kv_cache_spec_attn_group_iterator():
2850
+ for layer_name in group.layer_names:
2851
raw_tensor = kv_cache_raw_tensors[layer_name]
2852
num_blocks = (raw_tensor.numel() //
2853
kv_cache_spec.page_size_bytes)
2854
if isinstance(kv_cache_spec, AttentionSpec):
2861
- kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
2855
+
2856
+ kv_cache_shape = group.backend.get_kv_cache_shape(
2857
num_blocks, kv_cache_spec.block_size,
2858
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
2859
if kv_cache_shape[0] != num_blocks or kv_cache_shape[
@@ -2893,6 +2888,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2888
self.shared_kv_cache_layers,
2889
kv_cache_config.kv_cache_groups,
2890
kv_caches,
2891
+ self.attn_groups,
2892
)
2893
attn_layers = get_layers_from_vllm_config(self.vllm_config,
2894
Attention)
@@ -2958,9 +2954,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2954
continue
2955
2956
# TODO: Support other attention modules, e.g., cross-attention
2957
+ # TODO(lucas): move the attention specs into the model layers like
2958
+ # the attention backends
2959
if attn_module.attn_type == AttentionType.DECODER:
2962
- use_local_attention = (self.attention_chunk_size is not None
2963
- and attn_module.use_irope)
2960
if attn_module.sliding_window is not None:
2961
kv_cache_spec[layer_name] = SlidingWindowSpec(
2962
block_size=block_size,
@@ -2969,10 +2965,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2965
dtype=self.kv_cache_dtype,
2966
sliding_window=attn_module.sliding_window,
2967
use_mla=use_mla)
2972
- assert not use_local_attention, (
2973
- "attention module can not be with ",
2974
- "both local attention and sliding window")
2975
- elif use_local_attention:
2968
+ elif self.attention_chunk_size is not None \
2969
+ and isinstance(attn_module, ChunkedLocalAttention):
2970
kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(
2971
block_size=block_size,
2972
num_kv_heads=attn_module.num_kv_heads,
@@ -3043,7 +3037,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3037
3038
# Use the first attention metadata builder
3039
# to create encoder attention metadata
3046
- builder = self.attn_metadata_builders[0]
3040
+ builder = self.attn_groups[0][0].metadata_builder
3041
3042
dummy_block_table = torch.zeros((num_reqs, 1),
3043
dtype=torch.int32,
vllm/v1/worker/tpu_model_runner.py
+3
-2
@@ -15,8 +15,9 @@ import torch_xla.distributed.spmd as xs
15
import torch_xla.runtime as xr
16
17
import vllm.envs as envs
18
+from vllm.attention import Attention
19
from vllm.attention.backends.abstract import AttentionType
19
-from vllm.attention.layer import Attention
20
+from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
21
from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
22
from vllm.config import (ParallelConfig, VllmConfig,
23
get_layers_from_vllm_config, update_config)
@@ -518,7 +519,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
519
continue
520
521
if attn_module.attn_type == AttentionType.DECODER:
521
- if attn_module.use_irope:
522
+ if isinstance(attn_module, ChunkedLocalAttention):
523
logger.warning_once(
524
"Using irope in Pallas is not supported yet, it "
525
"will fall back to global attention for long context.")
vllm/v1/worker/utils.py
+21
@@ -1,14 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections import defaultdict
4
+from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Optional
6
7
import torch
8
9
+from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.config import ModelConfig, SchedulerConfig
11
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
12
from vllm.model_executor.models.utils import extract_layer_index
13
from vllm.multimodal.registry import MultiModalRegistry
14
+from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
15
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
16
from vllm.v1.kv_cache_interface import KVCacheGroupSpec
17
@@ -122,6 +125,13 @@ class MultiModalBudget:
125
return max_items_per_prompt, max_items_per_batch
126
127
128
+@dataclass
129
+class AttentionGroup:
130
+ backend: type[AttentionBackend]
131
+ metadata_builder: AttentionMetadataBuilder
132
+ layer_names: list[str]
133
+
134
+
135
def sanity_check_mm_encoder_outputs(
136
mm_embeddings: MultiModalEmbeddings,
137
expected_num_items: int,
@@ -196,6 +206,8 @@ def initialize_kv_cache_for_kv_sharing(
206
shared_kv_cache_layers: dict[str, str],
207
kv_cache_groups: list[KVCacheGroupSpec],
208
kv_caches: dict[str, torch.Tensor],
209
+ # Optional for now to avoid breaking TPU
210
+ attn_groups: Optional[list[list[AttentionGroup]]] = None,
211
) -> None:
212
"""
213
Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
@@ -225,6 +237,15 @@ def initialize_kv_cache_for_kv_sharing(
237
group_idx = layer_to_kv_cache_group_idx[target_layer_name]
238
kv_cache_groups[group_idx].layer_names.append(layer_name)
239
240
+ if attn_groups is not None:
241
+ assert len(attn_groups[group_idx]) == 1, (
242
+ "Only one attention group per KV cache group is supported "
243
+ "for KV-cache sharing for now.")
244
+ # TODO(lucas): I think in the future the layers that re-use a
245
+ # KV cache will be in a different attention group so we can
246
+ # remove this code from here.
247
+ attn_groups[group_idx][0].layer_names.append(layer_name)
248
+
249
250
def bind_kv_cache(
251
kv_caches: dict[str, torch.Tensor],