[Attention]: Use KVCacheSpec for AttentionMetadataBuilder type hints (#50148)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
Martin Hickey committed
Jul 31, 2026 at 12:20 UTC
b2fb83e7ffbc30a1aa4667b1dad7ca3e2c342bcf
14 files changed
+34
-33
vllm/model_executor/layers/attention/chunked_local_attention.py
+1
-2
@@ -20,7 +20,6 @@ from vllm.v1.attention.backends.utils import (
20
)
21
from vllm.v1.attention.selector import get_attn_backend
22
from vllm.v1.kv_cache_interface import (
23
- AttentionSpec,
23
ChunkedLocalAttentionSpec,
24
KVCacheSpec,
25
get_kv_quant_mode,
@@ -42,7 +41,7 @@ def create_chunked_local_attention_backend(
41
def get_cudagraph_support(
42
cls: type["AttentionMetadataBuilder"],
43
vllm_config: VllmConfig,
45
- kv_cache_spec: AttentionSpec,
44
+ kv_cache_spec: KVCacheSpec,
45
) -> AttentionCGSupport:
46
# Explicit override in case the underlying builder specialized this getter.
47
# @override omitted only because of mypy limitation due to type variable.
vllm/model_executor/layers/attention/mla_attention.py
+2
@@ -1722,6 +1722,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1722
understand this class
1723
"""
1724
1725
+ kv_cache_spec: AttentionSpec
1726
+
1727
# Defines the level of query length support for this backend.
1728
# - SINGLE_ONLY: Only single-token queries (no spec decode support)
1729
# - UNIFORM: Supports uniform multi-token queries (spec decode with uniform lengths)
vllm/v1/attention/backend.py
+3
-3
@@ -25,7 +25,7 @@ if TYPE_CHECKING:
25
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
26
from vllm.platforms.interface import DeviceCapability
27
from vllm.v1.attention.backends.utils import KVCacheLayoutType
28
- from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode
28
+ from vllm.v1.kv_cache_interface import KVCacheSpec, KVQuantMode
29
30
from vllm.v1.kv_cache_interface import get_kv_quant_mode
31
@@ -635,7 +635,7 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
635
@abstractmethod
636
def __init__(
637
self,
638
- kv_cache_spec: "AttentionSpec",
638
+ kv_cache_spec: "KVCacheSpec",
639
layer_names: list[str],
640
vllm_config: "VllmConfig",
641
device: torch.device,
@@ -649,7 +649,7 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
649
def get_cudagraph_support(
650
cls: type["AttentionMetadataBuilder"],
651
vllm_config: "VllmConfig",
652
- kv_cache_spec: "AttentionSpec",
652
+ kv_cache_spec: "KVCacheSpec",
653
) -> AttentionCGSupport:
654
"""Get the cudagraph support level of this builder class."""
655
return cls._cudagraph_support
vllm/v1/attention/backends/flash_attn.py
+2
-2
@@ -58,7 +58,7 @@ from vllm.v1.attention.backend import (
58
CommonAttentionMetadata,
59
)
60
from vllm.v1.attention.backends.utils import get_kv_cache_layout
61
-from vllm.v1.kv_cache_interface import AttentionSpec
61
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
62
from vllm.v1.worker.cp_utils import (
63
run_split_fa2_dcp_context_attention,
64
should_skip_dcp_context_attention,
@@ -360,7 +360,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
360
def get_cudagraph_support(
361
cls,
362
vllm_config: "VllmConfig",
363
- kv_cache_spec: "AttentionSpec",
363
+ kv_cache_spec: "KVCacheSpec",
364
) -> AttentionCGSupport:
365
return cls._cudagraph_support
366
vllm/v1/attention/backends/flashinfer.py
+3
-1
@@ -76,6 +76,7 @@ from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
76
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
77
from vllm.v1.kv_cache_interface import (
78
AttentionSpec,
79
+ KVCacheSpec,
80
KVQuantMode,
81
UniformTypeKVCacheSpecs,
82
)
@@ -617,6 +618,7 @@ class FlashInferMetadata:
618
619
620
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
621
+ kv_cache_spec: AttentionSpec
622
reorder_batch_threshold: int = 1
623
624
def __init__(
@@ -893,7 +895,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
895
def get_cudagraph_support(
896
cls: type["FlashInferMetadataBuilder"],
897
vllm_config: VllmConfig,
896
- kv_cache_spec: AttentionSpec,
898
+ kv_cache_spec: KVCacheSpec,
899
) -> AttentionCGSupport:
900
"""Get the cudagraph support level for FlashInfer attention.
901
vllm/v1/attention/backends/gdn_attn.py
+3
-3
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (
21
mamba_get_block_table_tensor,
22
split_decodes_and_prefills,
23
)
24
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
24
+from vllm.v1.kv_cache_interface import MambaSpec
25
26
27
class GDNAttentionBackend(AttentionBackend):
@@ -80,18 +80,18 @@ class GDNAttentionMetadata:
80
81
82
class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]):
83
+ kv_cache_spec: MambaSpec
84
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
85
86
reorder_batch_threshold: int = 1
87
88
def __init__(
89
self,
89
- kv_cache_spec: AttentionSpec,
90
+ kv_cache_spec: MambaSpec,
91
layer_names: list[str],
92
vllm_config: VllmConfig,
93
device: torch.device,
94
):
94
- assert isinstance(kv_cache_spec, MambaSpec)
95
self.vllm_config = vllm_config
96
self.compilation_config = vllm_config.compilation_config
97
self.speculative_config = vllm_config.speculative_config
vllm/v1/attention/backends/hpc_attn.py
+2
-2
@@ -33,7 +33,7 @@ from vllm.v1.attention.backends.utils import (
33
KVCacheLayoutType,
34
split_decodes_and_prefills,
35
)
36
-from vllm.v1.kv_cache_interface import AttentionSpec
36
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
37
38
logger = init_logger(__name__)
39
@@ -158,7 +158,7 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
158
def get_cudagraph_support(
159
cls: type["HpcAttnMetadataBuilder"],
160
vllm_config: VllmConfig,
161
- kv_cache_spec: AttentionSpec,
161
+ kv_cache_spec: KVCacheSpec,
162
) -> AttentionCGSupport:
163
spec_config = vllm_config.speculative_config
164
if (
vllm/v1/attention/backends/linear_attn.py
+5
-5
@@ -16,7 +16,7 @@ from vllm.v1.attention.backends.utils import (
16
mamba_get_block_table_tensor,
17
split_decodes_and_prefills,
18
)
19
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
19
+from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
20
21
22
class LinearAttentionBackend(AttentionBackend):
@@ -46,19 +46,19 @@ class LinearAttentionMetadata:
46
47
48
class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMetadata]):
49
+ kv_cache_spec: MambaSpec
50
reorder_batch_threshold: int = 1
51
52
_cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
53
54
def __init__(
55
self,
55
- kv_cache_spec: AttentionSpec,
56
+ kv_cache_spec: MambaSpec,
57
layer_names: list[str],
58
vllm_config: VllmConfig,
59
device: torch.device,
60
):
61
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
61
- assert isinstance(kv_cache_spec, MambaSpec)
62
63
def build(
64
self,
@@ -120,13 +120,13 @@ class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder):
120
def get_cudagraph_support(
121
cls,
122
vllm_config: VllmConfig,
123
- kv_cache_spec: AttentionSpec,
123
+ kv_cache_spec: KVCacheSpec,
124
) -> AttentionCGSupport:
125
return AttentionCGSupport.UNIFORM_BATCH
126
127
def __init__(
128
self,
129
- kv_cache_spec: AttentionSpec,
129
+ kv_cache_spec: MambaSpec,
130
layer_names: list[str],
131
vllm_config: VllmConfig,
132
device: torch.device,
vllm/v1/attention/backends/mamba2_attn.py
+2
-2
@@ -16,7 +16,7 @@ from vllm.v1.attention.backends.mamba_attn import (
16
BaseMambaAttentionMetadata,
17
BaseMambaAttentionMetadataBuilder,
18
)
19
-from vllm.v1.kv_cache_interface import AttentionSpec
19
+from vllm.v1.kv_cache_interface import MambaSpec
20
21
22
def compute_varlen_chunk_metadata(
@@ -118,7 +118,7 @@ class Mamba2AttentionMetadataBuilder(
118
119
def __init__(
120
self,
121
- kv_cache_spec: AttentionSpec,
121
+ kv_cache_spec: MambaSpec,
122
layer_names: list[str],
123
vllm_config: VllmConfig,
124
device: torch.device,
vllm/v1/attention/backends/mamba_attn.py
+3
-3
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (
21
mamba_get_block_table_tensor,
22
split_decodes_and_prefills,
23
)
24
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
24
+from vllm.v1.kv_cache_interface import MambaSpec
25
26
M = TypeVar("M", bound="BaseMambaAttentionMetadata")
27
@@ -83,6 +83,7 @@ class BaseMambaAttentionMetadata:
83
84
85
class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
86
+ kv_cache_spec: MambaSpec
87
metadata_cls: type[M]
88
reorder_batch_threshold: int = 1
89
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
@@ -92,7 +93,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
93
94
def __init__(
95
self,
95
- kv_cache_spec: AttentionSpec,
96
+ kv_cache_spec: MambaSpec,
97
layer_names: list[str],
98
vllm_config: VllmConfig,
99
device: torch.device,
@@ -107,7 +108,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
108
self.use_replayssm = vllm_config.cache_config.use_replayssm
109
self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
110
110
- assert isinstance(kv_cache_spec, MambaSpec)
111
scheduler_config = vllm_config.scheduler_config
112
self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs
113
if self.compilation_config.max_cudagraph_capture_size is not None:
vllm/v1/attention/backends/mla/indexer.py
+2
-2
@@ -36,7 +36,7 @@ from vllm.v1.attention.backends.utils import (
36
get_dcp_local_seq_lens,
37
split_decodes_and_prefills,
38
)
39
-from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
39
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
40
from vllm.v1.worker.cp_utils import get_kv_cache_shard_count
41
42
logger = init_logger(__name__)
@@ -463,7 +463,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
463
def get_cudagraph_support(
464
cls,
465
vllm_config: VllmConfig,
466
- kv_cache_spec: AttentionSpec,
466
+ kv_cache_spec: KVCacheSpec,
467
) -> AttentionCGSupport:
468
return AttentionCGSupport.UNIFORM_BATCH
469
vllm/v1/attention/backends/turboquant_attn.py
+2
@@ -53,6 +53,7 @@ from vllm.v1.attention.ops.triton_turboquant_decode import (
53
triton_turboquant_decode_attention,
54
)
55
from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
56
+from vllm.v1.kv_cache_interface import AttentionSpec
57
from vllm.v1.worker.workspace import (
58
current_workspace_manager,
59
is_workspace_manager_initialized,
@@ -198,6 +199,7 @@ class TurboQuantMetadata(AttentionMetadata):
199
class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):
200
"""Builds TurboQuantMetadata from scheduler output."""
201
202
+ kv_cache_spec: AttentionSpec
203
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
204
205
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
vllm/v1/worker/gpu/attn_utils.py
+1
-1
@@ -169,7 +169,7 @@ def init_attn_backend(
169
# Check cudagraph support for the attention backend
170
cg_support = builder.get_cudagraph_support(
171
vllm_config,
172
- cast(AttentionSpec, group.kv_cache_spec),
172
+ group.kv_cache_spec,
173
)
174
if cg_support.value < min_cg_support.value:
175
min_cg_support = cg_support
vllm/v1/worker/gpu/model_states/encoder_only.py
+3
-7
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, cast
3
+from typing import Any
4
5
import torch
6
import torch.nn as nn
@@ -15,11 +15,7 @@ from vllm.v1.attention.backend import (
15
CommonAttentionMetadata,
16
)
17
from vllm.v1.core.sched.output import NewRequestData
18
-from vllm.v1.kv_cache_interface import (
19
- AttentionSpec,
20
- EncoderOnlyAttentionSpec,
21
- KVCacheConfig,
22
-)
18
+from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, KVCacheConfig
19
from vllm.v1.worker.gpu.input_batch import InputBatch
20
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
21
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
@@ -152,7 +148,7 @@ class EncoderOnlyModelState(DefaultModelState):
148
for group in self.encoder_attn_groups:
149
builder = group.get_metadata_builder(0)
150
cg_support = builder.get_cudagraph_support(
155
- self.vllm_config, cast(AttentionSpec, group.kv_cache_spec)
151
+ self.vllm_config, group.kv_cache_spec
152
)
153
if cg_support.value < support.value:
154
support = cg_support