| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | |
| 11 | from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config |
| 12 | from vllm.distributed.eplb.eplb_state import EplbState |
| 13 | from vllm.forward_context import set_forward_context |
| 14 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 15 | from vllm.model_executor.model_loader import get_model |
| 16 | from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata |
| 17 | from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher |
| 18 | from vllm.v1.utils import CpuGpuBuffer |
| 19 | from vllm.v1.worker.dp_utils import coordinate_batch_across_dp |
| 20 | from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 24 | |
| 25 | PADDING_SLOT_ID = -1 |
| 26 | |
| 27 | |
| 28 | class ExtractHiddenStatesProposer: |
| 29 | def __init__(self, vllm_config: VllmConfig, device): |
| 30 | assert vllm_config.speculative_config is not None |
| 31 | |
| 32 | self.num_speculative_tokens = ( |
| 33 | vllm_config.speculative_config.num_speculative_tokens |
| 34 | ) |
| 35 | assert self.num_speculative_tokens == 1 |
| 36 | if vllm_config.speculative_config.disable_padded_drafter_batch: |
| 37 | raise ValueError( |
| 38 | "disable_padded_drafter_batch is not supported with " |
| 39 | "extract_hidden_states method" |
| 40 | ) |
| 41 | self.vllm_config = vllm_config |
| 42 | self.device = device |
| 43 | self.dtype = vllm_config.model_config.dtype |
| 44 | self.dp_rank = vllm_config.parallel_config.data_parallel_rank |
| 45 | |
| 46 | self.eplb_state: EplbState | None = None |
| 47 | |
| 48 | # Model and attention layer tracking (initialized in load_model) |
| 49 | self.model: nn.Module | None = None |
| 50 | self.attn_layer_names: list[str] = [] |
| 51 | self.attn_metadata_builder: AttentionMetadataBuilder | None = None |
| 52 | self.kv_cache_gid: int = -1 |
| 53 | |
| 54 | # Maximum number of tokens for buffers |
| 55 | max_batch_size = vllm_config.scheduler_config.max_num_seqs |
| 56 | self.max_num_tokens = ( |
| 57 | vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size |
| 58 | ) |
| 59 | |
| 60 | self.backup_next_token_ids = CpuGpuBuffer( |
| 61 | max_batch_size, dtype=torch.int32, device=device |
| 62 | ) |
| 63 | |
| 64 | self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config |
| 65 | layer_ids = getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", None) |
| 66 | if not layer_ids: |
| 67 | raise ValueError( |
| 68 | "eagle_aux_hidden_state_layer_ids must be set in the draft " |
| 69 | "model config for extract_hidden_states method" |
| 70 | ) |
| 71 | self.num_hidden_states = len(layer_ids) |
| 72 | self.hidden_size = vllm_config.model_config.get_hidden_size() |
| 73 | self.hidden_states = torch.zeros( |
| 74 | (self.max_num_tokens, self.num_hidden_states, self.hidden_size), |
| 75 | dtype=self.dtype, |
| 76 | device=device, |
| 77 | ) |
| 78 | self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config) |
| 79 | |
| 80 | self._slot_mapping_buffer = torch.zeros( |
| 81 | self.max_num_tokens, dtype=torch.int64, device=device |
| 82 | ) |
| 83 | |
| 84 | def set_eplb_state(self, eplb_state: EplbState) -> None: |
| 85 | """Inject EPLB state after construction.""" |
| 86 | self.eplb_state = eplb_state |
| 87 | |
| 88 | def propose( |
| 89 | self, |
| 90 | num_speculative_tokens: int, |
| 91 | sampled_token_ids: torch.Tensor, |
| 92 | target_hidden_states: list[torch.Tensor], |
| 93 | common_attn_metadata: CommonAttentionMetadata, |
| 94 | slot_mappings: dict[str, torch.Tensor] |
| 95 | | list[dict[str, torch.Tensor]] |
| 96 | | None = None, |
| 97 | ) -> torch.Tensor: |
| 98 | """Propose draft tokens by calling the ExtractHiddenStatesModel model. |
| 99 | |
| 100 | The ExtractHiddenStatesModel caches the hidden states in the KV cache |
| 101 | without performing actual attention computation. This allows us to |
| 102 | extract and store hidden states for later use (e.g., KV transfer). |
| 103 | |
| 104 | This proposer doesn't actually perform speculation - it returns the |
| 105 | sampled tokens as "draft" tokens, ensuring they always verify (match). |
| 106 | The main purpose is to cache hidden states, not to speculate. |
| 107 | |
| 108 | Args: |
| 109 | num_speculative_tokens: Number of draft tokens to return per request |
| 110 | sampled_token_ids: Sampled token IDs from the target model |
| 111 | target_hidden_states: List of hidden state tensors from target model |
| 112 | (one per aux hidden state layer) |
| 113 | common_attn_metadata: Attention metadata |
| 114 | slot_mappings: Slot mappings for KV cache (unused, provided for |
| 115 | interface compatibility) |
| 116 | |
| 117 | Returns: |
| 118 | Tuple of: |
| 119 | - Draft tokens matching sampled tokens, shape [batch_size, 1] |
| 120 | - KV connector output (if KV transfer is active), else None |
| 121 | |
| 122 | """ |
| 123 | assert num_speculative_tokens == self.num_speculative_tokens |
| 124 | assert self.model is not None and isinstance(target_hidden_states, list) |
| 125 | |
| 126 | # target_hidden_states is a list of tensors (one per layer) |
| 127 | # Each tensor has shape [num_tokens, hidden_size] |
| 128 | # Stack to shape: [num_tokens, num_hidden_states, hidden_size] |
| 129 | stacked_hidden_states = torch.stack(target_hidden_states, dim=1) |
| 130 | num_tokens = stacked_hidden_states.shape[0] |
| 131 | |
| 132 | # Copy hidden states to buffer |
| 133 | self.hidden_states[:num_tokens] = stacked_hidden_states |
| 134 | |
| 135 | assert self.attn_metadata_builder is not None |
| 136 | attn_metadata = self.attn_metadata_builder.build_for_drafting( |
| 137 | common_attn_metadata=common_attn_metadata, draft_index=0 |
| 138 | ) |
| 139 | |
| 140 | # We assume all cache-only layers belong to the same KV cache group, |
| 141 | # thus using the same attention metadata. |
| 142 | per_layer_attn_metadata = {} |
| 143 | for layer_name in self.attn_layer_names: |
| 144 | per_layer_attn_metadata[layer_name] = attn_metadata |
| 145 | |
| 146 | cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( |
| 147 | self._determine_batch_execution_and_padding(num_tokens) |
| 148 | ) |
| 149 | if num_tokens_across_dp is not None: |
| 150 | num_tokens_across_dp[self.dp_rank] = num_input_tokens |
| 151 | |
| 152 | if self.eplb_state is not None: |
| 153 | assert self.vllm_config.speculative_config is not None |
| 154 | self.eplb_state.prepare_forward( |
| 155 | self.vllm_config.speculative_config.draft_model_config, |
| 156 | num_tokens, |
| 157 | ) |
| 158 | with set_forward_context( |
| 159 | per_layer_attn_metadata, |
| 160 | self.vllm_config, |
| 161 | num_tokens=num_input_tokens, |
| 162 | num_tokens_across_dp=num_tokens_across_dp, |
| 163 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 164 | slot_mapping=self._get_slot_mapping( |
| 165 | num_input_tokens, common_attn_metadata.slot_mapping |
| 166 | ), |
| 167 | ): |
| 168 | self.model( |
| 169 | hidden_states=self.hidden_states[:num_input_tokens], |
| 170 | ) |
| 171 | |
| 172 | # Return the sampled tokens as "draft" tokens |
| 173 | # Shape: [batch_size, 1] to match num_speculative_tokens=1 |
| 174 | # On decode steps with spec tokens, sampled_token_ids may have |
| 175 | # shape [batch_size, 2] (target + spec verification); slice to |
| 176 | # return only the target-sampled column. |
| 177 | return sampled_token_ids[:, :1] |
| 178 | |
| 179 | def _get_slot_mapping( |
| 180 | self, |
| 181 | num_tokens: int, |
| 182 | slot_mapping: torch.Tensor | None = None, |
| 183 | ) -> dict[str, torch.Tensor]: |
| 184 | """Return slot_mapping dict for cache-only attention layers. |
| 185 | |
| 186 | If slot_mapping is provided, copies it into the buffer first. |
| 187 | """ |
| 188 | if slot_mapping is not None: |
| 189 | num_actual = slot_mapping.shape[0] |
| 190 | self._slot_mapping_buffer[:num_actual].copy_(slot_mapping) |
| 191 | if num_tokens > num_actual: |
| 192 | self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID) |
| 193 | |
| 194 | view = self._slot_mapping_buffer[:num_tokens] |
| 195 | return {name: view for name in self.attn_layer_names} |
| 196 | |
| 197 | def _determine_batch_execution_and_padding( |
| 198 | self, |
| 199 | num_tokens: int, |
| 200 | use_cudagraphs: bool = True, |
| 201 | ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]: |
| 202 | cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch( |
| 203 | num_tokens, |
| 204 | valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None), |
| 205 | ) |
| 206 | num_tokens_padded = batch_desc.num_tokens |
| 207 | |
| 208 | # Extra coordination when running data-parallel since we need to |
| 209 | # coordinate across ranks |
| 210 | # TODO(Flechman): support DBO ubatching |
| 211 | should_ubatch, num_tokens_across_dp = False, None |
| 212 | if self.vllm_config.parallel_config.data_parallel_size > 1: |
| 213 | should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( |
| 214 | coordinate_batch_across_dp( |
| 215 | num_tokens_unpadded=num_tokens, |
| 216 | parallel_config=self.vllm_config.parallel_config, |
| 217 | allow_microbatching=False, |
| 218 | num_tokens_padded=num_tokens_padded, |
| 219 | cudagraph_mode=cudagraph_mode.value, |
| 220 | ) |
| 221 | ) |
| 222 | assert not should_ubatch, ( |
| 223 | "DBO ubatching not implemented for extract_hidden_states" |
| 224 | ) |
| 225 | |
| 226 | # Extract DP-synced values |
| 227 | if num_tokens_across_dp is not None: |
| 228 | dp_rank = self.dp_rank |
| 229 | num_tokens_padded = int(num_tokens_across_dp[dp_rank].item()) |
| 230 | # Re-dispatch with DP padding so we have the correct |
| 231 | # batch_descriptor |
| 232 | cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch( |
| 233 | num_tokens_padded, |
| 234 | valid_modes={CUDAGraphMode(synced_cudagraph_mode)}, |
| 235 | ) |
| 236 | # Assert to make sure the agreed upon token count is correct |
| 237 | # otherwise num_tokens_across_dp will no-longer be valid |
| 238 | assert batch_desc.num_tokens == num_tokens_padded |
| 239 | num_tokens_across_dp[dp_rank] = num_tokens_padded |
| 240 | |
| 241 | return cudagraph_mode, num_tokens_padded, num_tokens_across_dp |
| 242 | |
| 243 | def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None: |
| 244 | """Initialize cudagraph dispatcher keys. |
| 245 | |
| 246 | Only supports PIECEWISE cudagraphs (via mixed_mode). |
| 247 | Should be called after adjust_cudagraph_sizes_for_spec_decode. |
| 248 | """ |
| 249 | assert self.vllm_config.speculative_config is not None |
| 250 | if ( |
| 251 | not self.vllm_config.speculative_config.enforce_eager |
| 252 | and cudagraph_mode.mixed_mode() |
| 253 | in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL] |
| 254 | ): |
| 255 | proposer_cudagraph_mode = CUDAGraphMode.PIECEWISE |
| 256 | else: |
| 257 | proposer_cudagraph_mode = CUDAGraphMode.NONE |
| 258 | |
| 259 | self.cudagraph_dispatcher.initialize_cudagraph_keys(proposer_cudagraph_mode) |
| 260 | |
| 261 | @torch.inference_mode() |
| 262 | def dummy_run( |
| 263 | self, |
| 264 | num_tokens: int, |
| 265 | use_cudagraphs: bool = True, |
| 266 | is_graph_capturing: bool = False, |
| 267 | slot_mappings: dict[str, torch.Tensor] | None = None, |
| 268 | ) -> None: |
| 269 | assert self.model is not None, "Model must be initialized before dummy_run" |
| 270 | cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( |
| 271 | self._determine_batch_execution_and_padding( |
| 272 | num_tokens, use_cudagraphs=use_cudagraphs |
| 273 | ) |
| 274 | ) |
| 275 | |
| 276 | if num_tokens_across_dp is not None: |
| 277 | num_tokens_across_dp[self.dp_rank] = num_input_tokens |
| 278 | |
| 279 | # Use our own slot mapping buffer during cudagraph capture. |
| 280 | if ( |
| 281 | self.attn_layer_names |
| 282 | and slot_mappings is not None |
| 283 | and self.attn_layer_names[0] in slot_mappings |
| 284 | ): |
| 285 | slot_mapping_dict = self._get_slot_mapping(num_input_tokens) |
| 286 | else: |
| 287 | slot_mapping_dict = slot_mappings or {} |
| 288 | |
| 289 | with set_forward_context( |
| 290 | None, |
| 291 | self.vllm_config, |
| 292 | num_tokens=num_input_tokens, |
| 293 | num_tokens_across_dp=num_tokens_across_dp, |
| 294 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 295 | slot_mapping=slot_mapping_dict, |
| 296 | ): |
| 297 | self.model( |
| 298 | hidden_states=self.hidden_states[:num_input_tokens], |
| 299 | ) |
| 300 | |
| 301 | def _build_attn_metadata_builder( |
| 302 | self, draft_attn_layers: dict[str, AttentionLayerBase] |
| 303 | ) -> AttentionMetadataBuilder: |
| 304 | """Build the attention metadata builder from draft attention layers.""" |
| 305 | if not draft_attn_layers: |
| 306 | raise ValueError("No attention layers found for ExtractHiddenStatesModel") |
| 307 | layer = next(iter(draft_attn_layers.values())) |
| 308 | attn_backend = layer.get_attn_backend() |
| 309 | return attn_backend.get_builder_cls()( |
| 310 | layer.get_kv_cache_spec(self.vllm_config), |
| 311 | self.attn_layer_names, |
| 312 | self.vllm_config, |
| 313 | self.device, |
| 314 | ) |
| 315 | |
| 316 | def prepare_next_token_ids_padded( |
| 317 | self, |
| 318 | sampled_token_ids: torch.Tensor, |
| 319 | requests: dict[str, CachedRequestState], |
| 320 | gpu_input_batch: InputBatch, |
| 321 | discard_request_mask: torch.Tensor, |
| 322 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 323 | """Prepare next token IDs for speculative decoding. |
| 324 | |
| 325 | Since num_speculative_tokens == 1, sampled_token_ids has shape |
| 326 | (batch_size, 1). For each request we either use the sampled token |
| 327 | (if valid and not discarded) or a backup token from the request state. |
| 328 | """ |
| 329 | # Precompute backup token IDs for discarded requests. |
| 330 | num_reqs = gpu_input_batch.num_reqs |
| 331 | for i in range(num_reqs): |
| 332 | self.backup_next_token_ids.np[i] = requests[ |
| 333 | gpu_input_batch.req_ids[i] |
| 334 | ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1) |
| 335 | self.backup_next_token_ids.copy_to_gpu(num_reqs) |
| 336 | backup_tokens_gpu = self.backup_next_token_ids.gpu[:num_reqs] |
| 337 | |
| 338 | assert discard_request_mask.dtype == torch.bool |
| 339 | |
| 340 | # With num_speculative_tokens == 1, there is exactly one token |
| 341 | sampled = sampled_token_ids[:, 0] |
| 342 | is_valid = (sampled >= 0) & (sampled < gpu_input_batch.vocab_size) |
| 343 | valid_sampled_tokens_count = is_valid.to(torch.int32) |
| 344 | |
| 345 | use_sampled = is_valid & ~discard_request_mask[:num_reqs] |
| 346 | next_token_ids = torch.where( |
| 347 | use_sampled, sampled.to(torch.int32), backup_tokens_gpu |
| 348 | ) |
| 349 | |
| 350 | return next_token_ids, valid_sampled_tokens_count |
| 351 | |
| 352 | def load_model(self, target_model: nn.Module) -> None: |
| 353 | """Load the ExtractHiddenStatesModel model. |
| 354 | |
| 355 | This method instantiates the ExtractHiddenStatesModel model which is used |
| 356 | to cache hidden states during speculative decoding. The model uses |
| 357 | cache-only attention (no computation, just caching KV states). |
| 358 | |
| 359 | Args: |
| 360 | target_model: The target model (passed for compatibility with |
| 361 | EagleProposer interface, but not used here) |
| 362 | |
| 363 | """ |
| 364 | # Get the target model's attention layers before loading draft model |
| 365 | target_attn_layer_names = set( |
| 366 | get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys() # type: ignore[type-abstract] |
| 367 | ) |
| 368 | |
| 369 | assert self.vllm_config.speculative_config is not None |
| 370 | draft_model_config = self.vllm_config.speculative_config.draft_model_config |
| 371 | from vllm.compilation.backends import set_model_tag |
| 372 | |
| 373 | with set_model_tag("extract_hidden_states"): |
| 374 | self.model = get_model( |
| 375 | vllm_config=self.vllm_config, model_config=draft_model_config |
| 376 | ) |
| 377 | |
| 378 | # Identify draft model's attention layers (difference from target) |
| 379 | all_attn_layers = get_layers_from_vllm_config( |
| 380 | self.vllm_config, |
| 381 | AttentionLayerBase, # type: ignore[type-abstract] |
| 382 | ) |
| 383 | draft_attn_layers = { |
| 384 | name: layer |
| 385 | for name, layer in all_attn_layers.items() |
| 386 | if name not in target_attn_layer_names |
| 387 | } |
| 388 | self.attn_layer_names = list(draft_attn_layers.keys()) |
| 389 | assert len(draft_attn_layers) == 1, ( |
| 390 | "ExtractHiddenStatesModel should have exactly one " |
| 391 | f"attention layer, found {len(draft_attn_layers)}" |
| 392 | ) |
| 393 | self.attn_metadata_builder = self._build_attn_metadata_builder( |
| 394 | draft_attn_layers |
| 395 | ) |
| 396 | |
| 397 | def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: |
| 398 | """Validate all drafting layers belong to the same KV cache group |
| 399 | and record the group index for common_attn_metadata selection.""" |
| 400 | assert len(self.attn_layer_names) == 1 |
| 401 | layer = self.attn_layer_names[0] |
| 402 | for gid, group in enumerate(kv_cache_config.kv_cache_groups): |
| 403 | if layer in group.layer_names: |
| 404 | self.kv_cache_gid = gid |
| 405 | return |
| 406 | raise ValueError(f"Cache-only layer {layer!r} not in any KV cache group") |