| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from functools import cache |
| 5 | from typing import TYPE_CHECKING, NamedTuple, cast, get_args |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | from vllm.config.cache import CacheDType |
| 11 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 12 | from vllm.v1.attention.backend import AttentionBackend, AttentionType |
| 13 | from vllm.v1.attention.backends.registry import ( |
| 14 | MambaAttentionBackendEnum, |
| 15 | ) |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from vllm.v1.kv_cache_interface import KVCacheSpecKind |
| 19 | |
| 20 | |
| 21 | class AttentionSelectorConfig(NamedTuple): |
| 22 | head_size: int |
| 23 | dtype: torch.dtype |
| 24 | kv_cache_dtype: CacheDType | None |
| 25 | block_size: int | None |
| 26 | use_mla: bool = False |
| 27 | has_sink: bool = False |
| 28 | use_sparse: bool = False |
| 29 | use_mm_prefix: bool = False |
| 30 | use_per_head_quant_scales: bool = False |
| 31 | attn_type: str = AttentionType.DECODER |
| 32 | has_sliding_window: bool = False |
| 33 | use_non_causal: bool = False |
| 34 | use_batch_invariant: bool = False |
| 35 | use_kv_connector: bool = False |
| 36 | use_pcp: bool = False |
| 37 | use_adaptive_verification: bool = False |
| 38 | use_dcp: bool = False |
| 39 | use_rswa: bool = False |
| 40 | |
| 41 | def __repr__(self): |
| 42 | return ( |
| 43 | f"AttentionSelectorConfig(head_size={self.head_size}, " |
| 44 | f"dtype={self.dtype}, " |
| 45 | f"kv_cache_dtype={self.kv_cache_dtype}, " |
| 46 | f"block_size={self.block_size}, " |
| 47 | f"use_mla={self.use_mla}, " |
| 48 | f"has_sink={self.has_sink}, " |
| 49 | f"use_sparse={self.use_sparse}, " |
| 50 | f"use_mm_prefix={self.use_mm_prefix}, " |
| 51 | f"use_per_head_quant_scales={self.use_per_head_quant_scales}, " |
| 52 | f"attn_type={self.attn_type}, " |
| 53 | f"has_sliding_window={self.has_sliding_window}, " |
| 54 | f"use_non_causal={self.use_non_causal}, " |
| 55 | f"use_batch_invariant={self.use_batch_invariant}, " |
| 56 | f"use_kv_connector={self.use_kv_connector}, " |
| 57 | f"use_adaptive_verification={self.use_adaptive_verification}, " |
| 58 | f"use_pcp={self.use_pcp}, " |
| 59 | f"use_dcp={self.use_dcp}, " |
| 60 | f"use_rswa={self.use_rswa})" |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def get_attn_spec_kind( |
| 65 | use_mla: bool, |
| 66 | has_sliding_window: bool, |
| 67 | attn_type: str, |
| 68 | ) -> "KVCacheSpecKind": |
| 69 | """Derive the KV-cache group kind a layer belongs to from its signals. |
| 70 | |
| 71 | Mirrors ``get_kv_cache_spec_kind`` (which derives the kind from the |
| 72 | produced ``KVCacheSpec``) so users can target groups by kind when |
| 73 | setting ``AttentionConfig.backend_per_kind``. |
| 74 | |
| 75 | ``SINK_FULL_ATTENTION`` is intentionally not derived here: it is produced |
| 76 | only by the ``StaticSinkAttention`` layer, whereas a plain ``Attention`` |
| 77 | layer with attention sinks (e.g. gpt-oss) still yields a |
| 78 | ``FullAttentionSpec``/``SlidingWindowSpec``. Sinks therefore do not change |
| 79 | the kind. |
| 80 | |
| 81 | Args: |
| 82 | use_mla: Whether the layer uses multi-head latent attention. |
| 83 | has_sliding_window: Whether the layer applies a sliding window. |
| 84 | attn_type: The layer's ``AttentionType``. |
| 85 | |
| 86 | Returns: |
| 87 | The ``KVCacheSpecKind`` the layer maps to. |
| 88 | |
| 89 | """ |
| 90 | from vllm.v1.kv_cache_interface import KVCacheSpecKind |
| 91 | |
| 92 | if attn_type == AttentionType.ENCODER_ONLY: |
| 93 | return KVCacheSpecKind.ENCODER_ONLY_ATTENTION |
| 94 | if attn_type == AttentionType.ENCODER_DECODER: |
| 95 | return KVCacheSpecKind.CROSS_ATTENTION |
| 96 | if use_mla: |
| 97 | if has_sliding_window: |
| 98 | return KVCacheSpecKind.SLIDING_WINDOW_MLA |
| 99 | return KVCacheSpecKind.MLA_ATTENTION |
| 100 | if has_sliding_window: |
| 101 | return KVCacheSpecKind.SLIDING_WINDOW |
| 102 | return KVCacheSpecKind.FULL_ATTENTION |
| 103 | |
| 104 | |
| 105 | def get_attn_backend( |
| 106 | head_size: int, |
| 107 | dtype: torch.dtype, |
| 108 | kv_cache_dtype: str | None, |
| 109 | use_mla: bool = False, |
| 110 | has_sink: bool = False, |
| 111 | use_sparse: bool = False, |
| 112 | use_mm_prefix: bool = False, |
| 113 | use_per_head_quant_scales: bool = False, |
| 114 | attn_type: str | None = None, |
| 115 | num_heads: int | None = None, |
| 116 | has_sliding_window: bool = False, |
| 117 | ) -> type[AttentionBackend]: |
| 118 | """Selects which attention backend to use and lazily imports it.""" |
| 119 | if kv_cache_dtype is not None: |
| 120 | valid_cache_dtypes = get_args(CacheDType) |
| 121 | assert kv_cache_dtype in valid_cache_dtypes, ( |
| 122 | f"Invalid kv_cache_dtype: {kv_cache_dtype}. " |
| 123 | f"Valid values are: {valid_cache_dtypes}" |
| 124 | ) |
| 125 | |
| 126 | from vllm.config import get_current_vllm_config |
| 127 | |
| 128 | vllm_config = get_current_vllm_config() |
| 129 | |
| 130 | cache_config = vllm_config.cache_config |
| 131 | block_size: int | None |
| 132 | if cache_config is not None and cache_config.user_specified_block_size: |
| 133 | block_size = cache_config.block_size |
| 134 | else: |
| 135 | block_size = None |
| 136 | |
| 137 | kv_transfer_config = vllm_config.kv_transfer_config |
| 138 | use_kv_connector = ( |
| 139 | kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance |
| 140 | ) |
| 141 | |
| 142 | speculative_config = vllm_config.speculative_config |
| 143 | use_adaptive_verification = ( |
| 144 | speculative_config is not None |
| 145 | and speculative_config.enable_adaptive_verification |
| 146 | ) |
| 147 | if use_adaptive_verification: |
| 148 | from vllm.compilation.backends import model_tag |
| 149 | |
| 150 | # The drafter always runs full-length blocks; only the verifier sees |
| 151 | # the trimmed, device-decided query lengths. |
| 152 | use_adaptive_verification = model_tag != "dspark_head" |
| 153 | |
| 154 | attn_type = attn_type or AttentionType.DECODER |
| 155 | attn_selector_config = AttentionSelectorConfig( |
| 156 | head_size=head_size, |
| 157 | dtype=dtype, |
| 158 | kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype), |
| 159 | block_size=block_size, |
| 160 | use_mla=use_mla, |
| 161 | has_sink=has_sink, |
| 162 | use_sparse=use_sparse, |
| 163 | use_mm_prefix=use_mm_prefix, |
| 164 | use_per_head_quant_scales=use_per_head_quant_scales, |
| 165 | attn_type=attn_type, |
| 166 | has_sliding_window=has_sliding_window, |
| 167 | use_non_causal=vllm_config.attention_config.use_non_causal, |
| 168 | use_batch_invariant=envs.VLLM_BATCH_INVARIANT, |
| 169 | use_kv_connector=use_kv_connector, |
| 170 | use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1, |
| 171 | use_adaptive_verification=use_adaptive_verification, |
| 172 | use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1, |
| 173 | use_rswa=( |
| 174 | vllm_config.model_config is not None |
| 175 | and vllm_config.model_config.rswa_window is not None |
| 176 | ), |
| 177 | ) |
| 178 | |
| 179 | # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over |
| 180 | # the global backend; kinds not present in the map fall back to it. |
| 181 | attention_config = vllm_config.attention_config |
| 182 | backend = attention_config.backend |
| 183 | if attention_config.backend_per_kind: |
| 184 | kind = get_attn_spec_kind( |
| 185 | use_mla=use_mla, |
| 186 | has_sliding_window=has_sliding_window, |
| 187 | attn_type=attn_type, |
| 188 | ) |
| 189 | backend = attention_config.backend_per_kind.get(kind.value, backend) |
| 190 | |
| 191 | # The KV cache layout is resolved across all of the model's backends at once |
| 192 | # in get_kv_cache_spec(); a single selection cannot see its peers. |
| 193 | return _cached_get_attn_backend( |
| 194 | backend=backend, |
| 195 | attn_selector_config=attn_selector_config, |
| 196 | num_heads=num_heads, |
| 197 | _run_kv_cache_dtype=( |
| 198 | cache_config.cache_dtype if cache_config is not None else None |
| 199 | ), |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | @cache |
| 204 | def _cached_get_attn_backend( |
| 205 | backend, |
| 206 | attn_selector_config: AttentionSelectorConfig, |
| 207 | num_heads: int | None = None, |
| 208 | *, |
| 209 | _run_kv_cache_dtype: CacheDType | None = None, |
| 210 | ) -> type[AttentionBackend]: |
| 211 | # Some platform selectors inspect the run-wide cache dtype. Keep it in the |
| 212 | # cache key even though they read the value from the current config. |
| 213 | from vllm.platforms import current_platform |
| 214 | |
| 215 | attention_cls = current_platform.get_attn_backend_cls( |
| 216 | backend, |
| 217 | attn_selector_config=attn_selector_config, |
| 218 | num_heads=num_heads, |
| 219 | ) |
| 220 | if not attention_cls: |
| 221 | raise ValueError( |
| 222 | f"Invalid attention backend for {current_platform.device_name}" |
| 223 | ) |
| 224 | backend = resolve_obj_by_qualname(attention_cls) |
| 225 | return backend |
| 226 | |
| 227 | |
| 228 | def get_mamba_attn_backend( |
| 229 | mamba_type: MambaAttentionBackendEnum, |
| 230 | ) -> type[AttentionBackend]: |
| 231 | """Select which mamba attention backend to use and lazily import it.""" |
| 232 | return _cached_get_mamba_attn_backend(mamba_type) |
| 233 | |
| 234 | |
| 235 | @cache |
| 236 | def _cached_get_mamba_attn_backend( |
| 237 | mamba_type: MambaAttentionBackendEnum, |
| 238 | ) -> type[AttentionBackend]: |
| 239 | assert mamba_type and isinstance(mamba_type, MambaAttentionBackendEnum) |
| 240 | |
| 241 | mamba_attn_backend = mamba_type.get_class() |
| 242 | if envs.VLLM_BATCH_INVARIANT and not mamba_attn_backend.supports_batch_invariance(): |
| 243 | raise RuntimeError( |
| 244 | "VLLM batch_invariant mode is not supported for " |
| 245 | f"{mamba_attn_backend.get_name()}." |
| 246 | ) |
| 247 | return mamba_attn_backend |