| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Hidden States Extractor Model. |
| 5 | |
| 6 | This model extracts and caches hidden states from the target model |
| 7 | without performing actual token generation. It's used with the |
| 8 | extract_hidden_states speculative decoding method. |
| 9 | """ |
| 10 | |
| 11 | from collections.abc import Iterable |
| 12 | from dataclasses import replace |
| 13 | from typing import ClassVar |
| 14 | |
| 15 | import torch |
| 16 | import torch.nn as nn |
| 17 | |
| 18 | from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config |
| 19 | from vllm.config.cache import CacheDType |
| 20 | from vllm.forward_context import get_forward_context |
| 21 | from vllm.model_executor.layers.attention.attention import set_default_quant_scales |
| 22 | from vllm.model_executor.layers.attention.kv_transfer_utils import ( |
| 23 | maybe_transfer_kv_layer, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 26 | from vllm.model_executor.models.utils import maybe_prefix |
| 27 | from vllm.utils.torch_utils import is_quantized_kv_cache, kv_cache_dtype_str_to_dtype |
| 28 | from vllm.v1.attention.backend import ( |
| 29 | AttentionBackend, |
| 30 | AttentionImpl, |
| 31 | AttentionMetadataBuilder, |
| 32 | AttentionType, |
| 33 | CommonAttentionMetadata, |
| 34 | ) |
| 35 | from vllm.v1.kv_cache_interface import ( |
| 36 | AttentionSpec, |
| 37 | HiddenStateCacheSpec, |
| 38 | KVCacheSpec, |
| 39 | ) |
| 40 | |
| 41 | ########## Custom Ops ######## |
| 42 | |
| 43 | |
| 44 | def unified_kv_cache_update( |
| 45 | to_cache: torch.Tensor, |
| 46 | layer_name: str, |
| 47 | ) -> torch.Tensor: |
| 48 | """Returns a dummy that is passed to unified_attention to signal a side effect and |
| 49 | the data dependency between them to ensure torch.compile preserves ordering. |
| 50 | """ |
| 51 | forward_context = get_forward_context() |
| 52 | attn_layer = forward_context.no_compile_layers[layer_name] |
| 53 | kv_cache = attn_layer.kv_cache |
| 54 | |
| 55 | slot_mapping = forward_context.slot_mapping |
| 56 | assert isinstance(slot_mapping, dict), ( |
| 57 | f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. " |
| 58 | ) |
| 59 | layer_slot_mapping = slot_mapping.get(layer_name) |
| 60 | if layer_slot_mapping is not None: |
| 61 | assert hasattr(attn_layer.impl, "do_kv_cache_update"), ( |
| 62 | f"{attn_layer.impl.__class__.__name__} does not support kv cache update" |
| 63 | ) |
| 64 | attn_layer.impl.do_kv_cache_update( |
| 65 | attn_layer, |
| 66 | to_cache, |
| 67 | kv_cache, |
| 68 | layer_slot_mapping, |
| 69 | ) |
| 70 | |
| 71 | return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype) |
| 72 | |
| 73 | |
| 74 | @maybe_transfer_kv_layer |
| 75 | def dummy_attention(layer_name, _placeholder): |
| 76 | # Note: layer_name arg required by @maybe_transfer_kv_layer |
| 77 | return _placeholder |
| 78 | |
| 79 | |
| 80 | def basic_cache( |
| 81 | to_cache: torch.Tensor, # shape: [seq_len, num_heads, head_size] |
| 82 | kv_cache: torch.Tensor, # shape: [num_blocks, num_heads, block_size, head_size] |
| 83 | slot_mapping: torch.Tensor, # shape: [seq_len] |
| 84 | ): |
| 85 | # Padding slots are -1; redirect them to the null block (block 0, never |
| 86 | # allocated to a request) so the scatter stays branch-free and sync-free. |
| 87 | block_size = kv_cache.shape[2] |
| 88 | slot_mapping = slot_mapping.clamp_min(0) |
| 89 | kv_cache[slot_mapping // block_size, :, slot_mapping % block_size] = to_cache |
| 90 | |
| 91 | |
| 92 | ######### CacheOnlyAttentionBackend ######## |
| 93 | |
| 94 | |
| 95 | class CacheOnlyAttentionBackend(AttentionBackend): |
| 96 | """Attention backend that only caches KV without computing attention.""" |
| 97 | |
| 98 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 99 | torch.float16, |
| 100 | torch.bfloat16, |
| 101 | torch.float32, |
| 102 | ] |
| 103 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 104 | "auto", |
| 105 | "bfloat16", |
| 106 | ] |
| 107 | forward_includes_kv_cache_update: bool = False |
| 108 | |
| 109 | @staticmethod |
| 110 | def get_name() -> str: |
| 111 | return "CACHE_ONLY_ATTN" |
| 112 | |
| 113 | @classmethod |
| 114 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 115 | return attn_type == AttentionType.DECODER |
| 116 | |
| 117 | @classmethod |
| 118 | def supports_mm_prefix(cls) -> bool: |
| 119 | return True |
| 120 | |
| 121 | @staticmethod |
| 122 | def get_impl_cls() -> type["CacheOnlyAttentionImpl"]: |
| 123 | return CacheOnlyAttentionImpl |
| 124 | |
| 125 | @staticmethod |
| 126 | def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]: |
| 127 | return CacheOnlyAttentionMetadataBuilder |
| 128 | |
| 129 | @staticmethod |
| 130 | def use_cascade_attention(*args, **kwargs) -> bool: |
| 131 | return False |
| 132 | |
| 133 | @classmethod |
| 134 | def get_supported_head_sizes(cls) -> list[int]: |
| 135 | return [] |
| 136 | |
| 137 | |
| 138 | class CacheOnlyAttentionMetadata: |
| 139 | def __init__(self, slot_mapping: torch.Tensor): |
| 140 | self.slot_mapping = slot_mapping |
| 141 | |
| 142 | |
| 143 | class CacheOnlyAttentionMetadataBuilder( |
| 144 | AttentionMetadataBuilder[CacheOnlyAttentionMetadata] |
| 145 | ): |
| 146 | def __init__( |
| 147 | self, |
| 148 | kv_cache_spec: AttentionSpec, |
| 149 | layer_names: list[str], |
| 150 | vllm_config: VllmConfig, |
| 151 | device: torch.device, |
| 152 | ): |
| 153 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 154 | |
| 155 | def build( |
| 156 | self, |
| 157 | common_prefix_len: int, |
| 158 | common_attn_metadata: CommonAttentionMetadata, |
| 159 | fast_build: bool = False, |
| 160 | ) -> CacheOnlyAttentionMetadata: |
| 161 | use_cascade = common_prefix_len > 0 |
| 162 | if use_cascade: |
| 163 | raise NotImplementedError( |
| 164 | "Cascade attention not supported by CacheOnlyAttention" |
| 165 | ) |
| 166 | causal = common_attn_metadata.causal |
| 167 | if not causal: |
| 168 | raise NotImplementedError( |
| 169 | "Non-causal attention not supported by CacheOnlyAttention" |
| 170 | ) |
| 171 | |
| 172 | return CacheOnlyAttentionMetadata( |
| 173 | slot_mapping=common_attn_metadata.slot_mapping, |
| 174 | ) |
| 175 | |
| 176 | |
| 177 | class CacheOnlyAttentionImpl(AttentionImpl): |
| 178 | """Attention implementation that only caches KV states.""" |
| 179 | |
| 180 | def __init__( |
| 181 | self, |
| 182 | num_heads: int, |
| 183 | head_size: int, |
| 184 | kv_cache_dtype: str, |
| 185 | kv_cache_torch_dtype: torch.dtype, |
| 186 | attn_type: AttentionType = AttentionType.DECODER, |
| 187 | ) -> None: |
| 188 | self.num_heads = num_heads |
| 189 | self.head_size = head_size |
| 190 | self.kv_cache_dtype = kv_cache_dtype |
| 191 | self.kv_cache_torch_dtype = kv_cache_torch_dtype |
| 192 | |
| 193 | if attn_type != AttentionType.DECODER: |
| 194 | raise NotImplementedError(f"Unsupported attention type: {attn_type}") |
| 195 | if is_quantized_kv_cache(kv_cache_dtype): |
| 196 | raise NotImplementedError("Quantized KV cache not supported") |
| 197 | |
| 198 | self.num_queries_per_kv = 1 |
| 199 | |
| 200 | def do_kv_cache_update( |
| 201 | self, |
| 202 | layer, |
| 203 | to_cache, |
| 204 | kv_cache, |
| 205 | slot_mapping, |
| 206 | ): |
| 207 | assert to_cache.dtype == self.kv_cache_torch_dtype, ( |
| 208 | f"Data to cache must be {self.kv_cache_torch_dtype}, got {to_cache.dtype}" |
| 209 | ) |
| 210 | assert kv_cache.dtype == self.kv_cache_torch_dtype, ( |
| 211 | f"KV cache must be {self.kv_cache_torch_dtype}, got {kv_cache.dtype}" |
| 212 | ) |
| 213 | |
| 214 | basic_cache(to_cache, kv_cache, slot_mapping) |
| 215 | |
| 216 | def forward(self, *args, **kwargs): |
| 217 | # Empty implementation of abstract method |
| 218 | pass |
| 219 | |
| 220 | |
| 221 | ############## CacheOnlyAttentionLayer (replaces Attention) ############ |
| 222 | |
| 223 | |
| 224 | class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase): |
| 225 | """Attention layer that only caches key/value states without computing attention.""" |
| 226 | |
| 227 | def __init__( |
| 228 | self, |
| 229 | num_heads: int, |
| 230 | head_size: int, |
| 231 | cache_config: CacheConfig | None = None, |
| 232 | prefix: str = "", |
| 233 | attn_type: AttentionType = AttentionType.DECODER, |
| 234 | ): |
| 235 | super().__init__() |
| 236 | |
| 237 | self.num_heads = num_heads |
| 238 | self.head_size = head_size |
| 239 | self.layer_name = prefix |
| 240 | |
| 241 | vllm_config = get_current_vllm_config() |
| 242 | |
| 243 | # KV cache configuration |
| 244 | cache_config = cache_config or vllm_config.cache_config |
| 245 | if cache_config is not None: |
| 246 | kv_cache_dtype = cache_config.cache_dtype |
| 247 | self.block_size = cache_config.block_size |
| 248 | else: |
| 249 | kv_cache_dtype = "auto" |
| 250 | self.block_size = 16 |
| 251 | |
| 252 | assert kv_cache_dtype in ["auto", "bfloat16", "float16"], ( |
| 253 | "CacheOnlyAttentionLayer doesn't currently support quantized kv cache but" |
| 254 | f"kv cache dtype was set to {kv_cache_dtype}" |
| 255 | ) |
| 256 | self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( |
| 257 | kv_cache_dtype, vllm_config.model_config |
| 258 | ) |
| 259 | |
| 260 | # Initialize KV cache quantization attributes |
| 261 | set_default_quant_scales(self, register_buffer=True) |
| 262 | |
| 263 | # Attention backend |
| 264 | self.attn_backend = CacheOnlyAttentionBackend |
| 265 | impl_cls = self.attn_backend.get_impl_cls() |
| 266 | self.impl = impl_cls( |
| 267 | num_heads, |
| 268 | head_size, |
| 269 | kv_cache_dtype, |
| 270 | self.kv_cache_torch_dtype, |
| 271 | attn_type, |
| 272 | ) |
| 273 | |
| 274 | assert not self.attn_backend.forward_includes_kv_cache_update, ( |
| 275 | "KV cache update should be independent of forward" |
| 276 | ) |
| 277 | |
| 278 | # Placeholder KV cache (replaced by bind_kv_cache) |
| 279 | self.kv_cache = torch.tensor([]) |
| 280 | |
| 281 | # Register in compilation context |
| 282 | compilation_config = vllm_config.compilation_config |
| 283 | if prefix in compilation_config.static_forward_context: |
| 284 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 285 | compilation_config.static_forward_context[prefix] = self |
| 286 | |
| 287 | def forward(self, to_cache: torch.Tensor) -> torch.Tensor: |
| 288 | """Cache hidden states as KV pairs without computing attention. |
| 289 | |
| 290 | Args: |
| 291 | to_cache: The tensor to insert into the kv cache. |
| 292 | shape [num_tokens, num_heads, head_size] |
| 293 | |
| 294 | Returns: |
| 295 | Dummy output tensor (not used) |
| 296 | |
| 297 | """ |
| 298 | # Note: we set num_heads to num_hidden_layers and |
| 299 | # head_size to hidden_size for hidden states storage |
| 300 | output = torch.empty(0, device=to_cache.device, dtype=to_cache.dtype) |
| 301 | |
| 302 | # Note: dummy_out is used to force torch.compile to preserve ordering between |
| 303 | # cache update and attention op (which triggers kv_connector transfer) |
| 304 | dummy_out = unified_kv_cache_update(to_cache, self.layer_name) |
| 305 | |
| 306 | # Triggers kv_connector transfer via decorator |
| 307 | _ = dummy_attention(self.layer_name, dummy_out) |
| 308 | |
| 309 | return output |
| 310 | |
| 311 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 312 | return self.attn_backend |
| 313 | |
| 314 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 315 | # Re-read block_size: hybrid models may bump it after __init__. |
| 316 | return HiddenStateCacheSpec( |
| 317 | block_size=vllm_config.cache_config.block_size, |
| 318 | num_kv_heads=self.num_heads, |
| 319 | head_size=self.head_size, |
| 320 | dtype=self.kv_cache_torch_dtype, |
| 321 | ) |
| 322 | |
| 323 | |
| 324 | ############ ExtractHiddenStatesModel definition ########## |
| 325 | |
| 326 | |
| 327 | class ExtractHiddenStatesModel(nn.Module): |
| 328 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 329 | super().__init__() |
| 330 | |
| 331 | self.vllm_config = vllm_config |
| 332 | speculative_config = vllm_config.speculative_config |
| 333 | assert speculative_config is not None |
| 334 | self.hf_config = speculative_config.draft_model_config.hf_config |
| 335 | self.hidden_size = vllm_config.model_config.get_hidden_size() |
| 336 | self.target_num_hidden_layers = ( |
| 337 | vllm_config.model_config.get_total_num_hidden_layers() |
| 338 | ) |
| 339 | self.num_hidden_states = len( |
| 340 | getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", []) |
| 341 | ) |
| 342 | |
| 343 | cache_config = vllm_config.cache_config |
| 344 | |
| 345 | # Hidden states dtype should be independent of KV cache dtype. |
| 346 | if cache_config is not None and is_quantized_kv_cache(cache_config.cache_dtype): |
| 347 | cache_config = replace(cache_config, cache_dtype="auto") |
| 348 | |
| 349 | # Create a single cache-only attention layer |
| 350 | # Note: We set num_heads <- self.num_hidden_states |
| 351 | # and head_size <- hidden_size so that we can insert |
| 352 | # the hidden states directly into the cache without |
| 353 | # reshaping |
| 354 | self.cache_only_layers = nn.ModuleDict( |
| 355 | { |
| 356 | str(self.target_num_hidden_layers): CacheOnlyAttentionLayer( |
| 357 | num_heads=self.num_hidden_states, |
| 358 | head_size=self.hidden_size, |
| 359 | cache_config=cache_config, |
| 360 | prefix=maybe_prefix( |
| 361 | prefix, f"cache_only_layers.{self.target_num_hidden_layers}" |
| 362 | ), |
| 363 | ) |
| 364 | } |
| 365 | ) |
| 366 | |
| 367 | def forward(self, hidden_states: torch.Tensor) -> None: |
| 368 | """Process and cache hidden states. |
| 369 | |
| 370 | Args: |
| 371 | hidden_states: Hidden states from target model |
| 372 | shape: [num_tokens, num_hidden_states, hidden_size] |
| 373 | |
| 374 | Returns: |
| 375 | Tuple of (dummy_output, dummy_output) - both unused |
| 376 | |
| 377 | """ |
| 378 | # Call dummy attention layer to cache hidden states |
| 379 | # Output is ignored - we only care about the KV cache side effects |
| 380 | _ = self.cache_only_layers[str(self.target_num_hidden_layers)](hidden_states) |
| 381 | |
| 382 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 383 | """No weights to load for this dummy model.""" |
| 384 | return set() |