| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from copy import copy |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig, get_layers_from_vllm_config, replace |
| 9 | from vllm.forward_context import set_forward_context |
| 10 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 11 | from vllm.model_executor.models.utils import get_draft_quant_config |
| 12 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 13 | from vllm.v1.kv_cache_interface import ( |
| 14 | KVCacheConfig, |
| 15 | KVCacheSpec, |
| 16 | UniformTypeKVCacheSpecs, |
| 17 | ) |
| 18 | from vllm.v1.sample.metadata import SamplingMetadata |
| 19 | from vllm.v1.spec_decode.eagle import EagleProposer |
| 20 | from vllm.v1.spec_decode.utils import PADDING_SLOT_ID |
| 21 | from vllm.v1.worker.utils import AttentionGroup |
| 22 | |
| 23 | |
| 24 | class Step3p5MTPProposer(EagleProposer): |
| 25 | """Step3.5 MTP proposer with per-layer draft-step selection.""" |
| 26 | |
| 27 | def __init__( |
| 28 | self, |
| 29 | vllm_config: VllmConfig, |
| 30 | device: torch.device, |
| 31 | runner=None, |
| 32 | ): |
| 33 | super().__init__(vllm_config, device, runner) |
| 34 | self._per_group_block_tables: dict[int, torch.Tensor] = {} |
| 35 | self._per_group_slot_mappings: dict[int, torch.Tensor] = {} |
| 36 | # Slot-mapping buffers for non-primary KV cache groups (the primary |
| 37 | # group reuses self._slot_mapping_buffer from the base class). |
| 38 | self._per_group_slot_mapping_buffers: dict[int, torch.Tensor] = {} |
| 39 | |
| 40 | def set_per_group_attn_metadata( |
| 41 | self, |
| 42 | gid: int, |
| 43 | block_table: torch.Tensor, |
| 44 | slot_mapping: torch.Tensor, |
| 45 | ) -> None: |
| 46 | self._per_group_block_tables[gid] = block_table |
| 47 | self._per_group_slot_mappings[gid] = slot_mapping |
| 48 | |
| 49 | def _slot_mapping_buffer_for(self, gid: int) -> torch.Tensor: |
| 50 | if gid == self.kv_cache_gid: |
| 51 | return self._slot_mapping_buffer |
| 52 | buf = self._per_group_slot_mapping_buffers.get(gid) |
| 53 | if buf is None: |
| 54 | buf = torch.zeros(self.max_positions, dtype=torch.int64, device=self.device) |
| 55 | self._per_group_slot_mapping_buffers[gid] = buf |
| 56 | return buf |
| 57 | |
| 58 | def _get_slot_mapping( |
| 59 | self, |
| 60 | num_tokens: int, |
| 61 | slot_mapping: torch.Tensor | None = None, |
| 62 | ) -> dict[str, torch.Tensor]: |
| 63 | """Per-layer slot_mapping with one buffer per KV cache group.""" |
| 64 | per_layer: dict[str, torch.Tensor] = {} |
| 65 | for attn_group in self.draft_attn_groups: |
| 66 | gid = attn_group.kv_cache_group_id |
| 67 | buf = self._slot_mapping_buffer_for(gid) |
| 68 | source = self._per_group_slot_mappings.get(gid, slot_mapping) |
| 69 | if source is not None and buf.data_ptr() != source.data_ptr(): |
| 70 | n = source.shape[0] |
| 71 | buf[:n].copy_(source) |
| 72 | if num_tokens > n: |
| 73 | buf[n:num_tokens].fill_(PADDING_SLOT_ID) |
| 74 | view = buf[:num_tokens] |
| 75 | for layer_name in attn_group.layer_names: |
| 76 | per_layer[layer_name] = view |
| 77 | return per_layer |
| 78 | |
| 79 | def _update_positions_dependent_metadata( |
| 80 | self, |
| 81 | positions: torch.Tensor, |
| 82 | common_attn_metadata: CommonAttentionMetadata, |
| 83 | batch_size: int, |
| 84 | input_batch_size: int, |
| 85 | block_size: int, |
| 86 | ) -> torch.Tensor: |
| 87 | old_positions_1d = positions[0] if self.uses_mrope else positions |
| 88 | positions = super()._update_positions_dependent_metadata( |
| 89 | positions, |
| 90 | common_attn_metadata, |
| 91 | batch_size, |
| 92 | input_batch_size, |
| 93 | block_size, |
| 94 | ) |
| 95 | # Parent already produced slot_mapping for the primary gid. |
| 96 | self._per_group_slot_mappings[self.kv_cache_gid] = ( |
| 97 | common_attn_metadata.slot_mapping |
| 98 | ) |
| 99 | # Recompute slot_mapping for the remaining gids using their own block tables. |
| 100 | new_positions_1d = positions[0] if self.uses_mrope else positions |
| 101 | exceeds = old_positions_1d + 1 >= self.max_model_len |
| 102 | for attn_group in self.draft_attn_groups: |
| 103 | gid = attn_group.kv_cache_group_id |
| 104 | if gid == self.kv_cache_gid: |
| 105 | continue |
| 106 | block_table = self._per_group_block_tables.get(gid) |
| 107 | if block_table is None: |
| 108 | continue |
| 109 | n_blocks = block_table.shape[1] |
| 110 | bn = new_positions_1d // block_size |
| 111 | bn.clamp_(max=n_blocks - 1) |
| 112 | bn = bn.to(torch.long) |
| 113 | block_ids = block_table[:batch_size].gather(1, bn.unsqueeze(1)).squeeze(1) |
| 114 | sm = block_ids * block_size + (new_positions_1d % block_size) |
| 115 | sm.masked_fill_(exceeds, PADDING_SLOT_ID) |
| 116 | buf = self._slot_mapping_buffer_for(gid) |
| 117 | buf[:batch_size].copy_(sm) |
| 118 | if input_batch_size > batch_size: |
| 119 | buf[batch_size:input_batch_size].fill_(PADDING_SLOT_ID) |
| 120 | self._per_group_slot_mappings[gid] = buf[:batch_size] |
| 121 | return positions |
| 122 | |
| 123 | def build_per_group_and_layer_attn_metadata( |
| 124 | self, |
| 125 | common_attn_metadata: CommonAttentionMetadata, |
| 126 | draft_index: int = 0, |
| 127 | ) -> tuple[list[object], dict[str, object]]: |
| 128 | per_group_attn_metadata: list[object] = [] |
| 129 | per_layer_attn_metadata: dict[str, object] = {} |
| 130 | # The proposer always works in unpadded shape. Per-group block tables |
| 131 | # registered via set_per_group_attn_metadata are stored at the model |
| 132 | # runner's padded shape; slice them to match cm's num_reqs. |
| 133 | num_reqs = common_attn_metadata.num_reqs |
| 134 | num_actual_tokens = common_attn_metadata.num_actual_tokens |
| 135 | for attn_group in self.draft_attn_groups: |
| 136 | gid = attn_group.kv_cache_group_id |
| 137 | if gid in self._per_group_block_tables: |
| 138 | cm = copy(common_attn_metadata) |
| 139 | cm.block_table_tensor = self._per_group_block_tables[gid][:num_reqs] |
| 140 | if gid in self._per_group_slot_mappings: |
| 141 | sm = self._per_group_slot_mappings[gid] |
| 142 | if sm.shape[0] >= num_actual_tokens: |
| 143 | sm = sm[:num_actual_tokens] |
| 144 | cm.slot_mapping = sm |
| 145 | else: |
| 146 | cm = common_attn_metadata |
| 147 | attn_metadata = attn_group.get_metadata_builder().build_for_drafting( |
| 148 | common_attn_metadata=cm, |
| 149 | draft_index=draft_index, |
| 150 | ) |
| 151 | per_group_attn_metadata.append(attn_metadata) |
| 152 | for layer_name in attn_group.layer_names: |
| 153 | per_layer_attn_metadata[layer_name] = attn_metadata |
| 154 | return per_group_attn_metadata, per_layer_attn_metadata |
| 155 | |
| 156 | def _maybe_share_lm_head(self, target_language_model: torch.nn.Module) -> None: |
| 157 | """Step3.5 MTP uses the lm_head stored in each MTP layer.""" |
| 158 | # The base MTP path shares target lm_head into shared_head.head. |
| 159 | # Step3.5 checkpoints carry per-MTP-layer shared_head weights. |
| 160 | return |
| 161 | |
| 162 | def _create_draft_vllm_config(self) -> VllmConfig: |
| 163 | base = super()._create_draft_vllm_config() |
| 164 | return replace( |
| 165 | base, |
| 166 | model_config=self.draft_model_config, |
| 167 | quant_config=get_draft_quant_config(base), |
| 168 | ) |
| 169 | |
| 170 | def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: |
| 171 | """Step3.5 MTP draft layers may span multiple KV cache groups.""" |
| 172 | return |
| 173 | |
| 174 | def initialize_attn_backend( |
| 175 | self, |
| 176 | kv_cache_config: KVCacheConfig, |
| 177 | kernel_block_sizes: list[int] | None = None, |
| 178 | ) -> None: |
| 179 | all_attn_layers = get_layers_from_vllm_config( |
| 180 | self.vllm_config, |
| 181 | AttentionLayerBase, # type: ignore[type-abstract] |
| 182 | ) |
| 183 | |
| 184 | layer_to_gid: dict[str, int] = {} |
| 185 | layer_to_spec: dict[str, KVCacheSpec] = {} |
| 186 | for gid, group in enumerate(kv_cache_config.kv_cache_groups): |
| 187 | group_spec = group.kv_cache_spec |
| 188 | for layer_name in group.layer_names: |
| 189 | layer_to_gid[layer_name] = gid |
| 190 | if isinstance(group_spec, UniformTypeKVCacheSpecs): |
| 191 | if layer_name in group_spec.kv_cache_specs: |
| 192 | layer_to_spec[layer_name] = group_spec.kv_cache_specs[ |
| 193 | layer_name |
| 194 | ] |
| 195 | else: |
| 196 | target_layer_name = getattr( |
| 197 | all_attn_layers.get(layer_name), |
| 198 | "kv_sharing_target_layer_name", |
| 199 | None, |
| 200 | ) |
| 201 | if ( |
| 202 | target_layer_name |
| 203 | and target_layer_name in group_spec.kv_cache_specs |
| 204 | ): |
| 205 | layer_to_spec[layer_name] = group_spec.kv_cache_specs[ |
| 206 | target_layer_name |
| 207 | ] |
| 208 | else: |
| 209 | layer_to_spec[layer_name] = group_spec |
| 210 | else: |
| 211 | layer_to_spec[layer_name] = group_spec |
| 212 | |
| 213 | attention_groups: dict[tuple[tuple[str, str], int], AttentionGroup] = {} |
| 214 | for layer_name in sorted(self._draft_attn_layer_names): |
| 215 | if layer_name not in layer_to_spec: |
| 216 | continue |
| 217 | attn_layer = all_attn_layers[layer_name] |
| 218 | attn_backend = attn_layer.get_attn_backend() |
| 219 | spec = layer_to_spec[layer_name] |
| 220 | gid = layer_to_gid[layer_name] |
| 221 | group_key = (attn_backend.full_cls_name(), gid) |
| 222 | |
| 223 | if group_key not in attention_groups: |
| 224 | kernel_block_size = ( |
| 225 | kernel_block_sizes[gid] |
| 226 | if kernel_block_sizes is not None and gid < len(kernel_block_sizes) |
| 227 | else None |
| 228 | ) |
| 229 | attn_group = AttentionGroup( |
| 230 | backend=attn_backend, |
| 231 | layer_names=[layer_name], |
| 232 | kv_cache_spec=spec, |
| 233 | kv_cache_group_id=gid, |
| 234 | ) |
| 235 | attn_group.create_metadata_builders( |
| 236 | self.vllm_config, |
| 237 | self.device, |
| 238 | kernel_block_size=kernel_block_size, |
| 239 | ) |
| 240 | attention_groups[group_key] = attn_group |
| 241 | else: |
| 242 | attention_groups[group_key].layer_names.append(layer_name) |
| 243 | |
| 244 | self.draft_attn_groups = list(attention_groups.values()) |
| 245 | if self.draft_attn_groups: |
| 246 | self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id |
| 247 | self.block_size = ( |
| 248 | self.draft_attn_groups[0] |
| 249 | .get_metadata_builder() |
| 250 | .kv_cache_spec.block_size |
| 251 | ) |
| 252 | else: |
| 253 | self.kv_cache_gid = 0 |
| 254 | self.block_size = kv_cache_config.kv_cache_groups[ |
| 255 | 0 |
| 256 | ].kv_cache_spec.block_size |
| 257 | |
| 258 | def _sample_draft_tokens_for_step( |
| 259 | self, |
| 260 | hidden_states: torch.Tensor, |
| 261 | sampling_metadata: SamplingMetadata, |
| 262 | spec_step_idx: int, |
| 263 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 264 | if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy: |
| 265 | if self.use_local_argmax_reduction: |
| 266 | return self.model.get_top_tokens(hidden_states), None |
| 267 | logits = self.model.compute_logits( |
| 268 | hidden_states, spec_step_idx=spec_step_idx |
| 269 | ) |
| 270 | return logits.argmax(dim=-1), None |
| 271 | |
| 272 | logits = self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx) |
| 273 | return self._sample_from_logits(logits, sampling_metadata) |
| 274 | |
| 275 | def propose( |
| 276 | self, |
| 277 | num_speculative_tokens: int, |
| 278 | target_token_ids: torch.Tensor, |
| 279 | target_positions: torch.Tensor, |
| 280 | target_hidden_states: torch.Tensor, |
| 281 | next_token_ids: torch.Tensor, |
| 282 | token_indices_to_sample: torch.Tensor | None, |
| 283 | common_attn_metadata: CommonAttentionMetadata, |
| 284 | sampling_metadata: SamplingMetadata, |
| 285 | mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, |
| 286 | num_rejected_tokens_gpu: torch.Tensor | None = None, |
| 287 | slot_mappings: dict[str, torch.Tensor] |
| 288 | | list[dict[str, torch.Tensor]] |
| 289 | | None = None, |
| 290 | ) -> torch.Tensor: |
| 291 | self.num_speculative_tokens = num_speculative_tokens |
| 292 | self._last_draft_probs = None |
| 293 | batch_size = common_attn_metadata.batch_size() |
| 294 | |
| 295 | num_tokens, token_indices_to_sample, common_attn_metadata = ( |
| 296 | self.set_inputs_first_pass( |
| 297 | target_token_ids=target_token_ids, |
| 298 | next_token_ids=next_token_ids, |
| 299 | target_positions=target_positions, |
| 300 | target_hidden_states=target_hidden_states, |
| 301 | token_indices_to_sample=token_indices_to_sample, |
| 302 | cad=common_attn_metadata, |
| 303 | num_rejected_tokens_gpu=num_rejected_tokens_gpu, |
| 304 | ) |
| 305 | ) |
| 306 | |
| 307 | per_group_attn_metadata, per_layer_attn_metadata = ( |
| 308 | self.build_per_group_and_layer_attn_metadata(common_attn_metadata) |
| 309 | ) |
| 310 | |
| 311 | cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( |
| 312 | self._determine_batch_execution_and_padding(num_tokens) |
| 313 | ) |
| 314 | |
| 315 | model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass( |
| 316 | num_tokens, num_input_tokens, mm_embed_inputs |
| 317 | ) |
| 318 | model_kwargs["spec_step_idx"] = 0 |
| 319 | |
| 320 | with set_forward_context( |
| 321 | per_layer_attn_metadata, |
| 322 | self.vllm_config, |
| 323 | num_tokens=num_input_tokens, |
| 324 | num_tokens_across_dp=num_tokens_across_dp, |
| 325 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 326 | slot_mapping=self._get_slot_mapping( |
| 327 | slot_mapping_size, common_attn_metadata.slot_mapping |
| 328 | ), |
| 329 | ): |
| 330 | ret_hidden_states = self.model(**model_kwargs) |
| 331 | if not self.model_returns_tuple(): |
| 332 | last_hidden_states = ret_hidden_states |
| 333 | hidden_states = last_hidden_states |
| 334 | else: |
| 335 | last_hidden_states, hidden_states = ret_hidden_states |
| 336 | |
| 337 | sample_hidden_states = last_hidden_states[token_indices_to_sample] |
| 338 | |
| 339 | if self.num_speculative_tokens == 1 or self.parallel_drafting: |
| 340 | draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( |
| 341 | sample_hidden_states, sampling_metadata, spec_step_idx=0 |
| 342 | ) |
| 343 | if draft_probs is not None: |
| 344 | self._last_draft_probs = draft_probs.view( |
| 345 | -1, self.num_speculative_tokens, draft_probs.shape[-1] |
| 346 | ).contiguous() |
| 347 | return draft_token_ids.view(-1, self.num_speculative_tokens) |
| 348 | |
| 349 | if self.uses_mrope: |
| 350 | positions = self.mrope_positions[:, token_indices_to_sample] |
| 351 | else: |
| 352 | positions = self.positions[token_indices_to_sample] |
| 353 | hidden_states = hidden_states[token_indices_to_sample] |
| 354 | |
| 355 | if self.constant_draft_positions: |
| 356 | self.positions[:batch_size] = positions |
| 357 | |
| 358 | draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( |
| 359 | sample_hidden_states, sampling_metadata, spec_step_idx=0 |
| 360 | ) |
| 361 | draft_probs_list = None if draft_probs is None else [draft_probs] |
| 362 | |
| 363 | if self.allowed_attn_types is not None: |
| 364 | for group_md in per_group_attn_metadata: |
| 365 | if not isinstance(group_md, self.allowed_attn_types): |
| 366 | raise ValueError( |
| 367 | f"Unsupported attention metadata type for speculative " |
| 368 | "decoding with num_speculative_tokens > 1: " |
| 369 | f"{type(group_md)}. Supported types are: " |
| 370 | f"{self.allowed_attn_types}" |
| 371 | ) |
| 372 | |
| 373 | draft_token_ids_list = [draft_token_ids] |
| 374 | |
| 375 | cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = ( |
| 376 | self._determine_batch_execution_and_padding(batch_size) |
| 377 | ) |
| 378 | |
| 379 | common_attn_metadata.num_actual_tokens = batch_size |
| 380 | common_attn_metadata.max_query_len = 1 |
| 381 | common_attn_metadata.query_start_loc = self.arange[: batch_size + 1] |
| 382 | common_attn_metadata.query_start_loc_cpu = torch.from_numpy( |
| 383 | self.token_arange_np[: batch_size + 1] |
| 384 | ).clone() |
| 385 | |
| 386 | if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None: |
| 387 | common_attn_metadata.seq_lens -= num_rejected_tokens_gpu |
| 388 | |
| 389 | block_size = self.block_size |
| 390 | assert block_size > 0, "block_size has not been initialized." |
| 391 | for token_index in range(self.num_speculative_tokens - 1): |
| 392 | spec_step_idx = token_index + 1 |
| 393 | input_ids = draft_token_ids_list[-1].int() |
| 394 | |
| 395 | if not self.constant_draft_positions: |
| 396 | positions = self._update_positions_dependent_metadata( |
| 397 | positions, |
| 398 | common_attn_metadata, |
| 399 | batch_size, |
| 400 | input_batch_size, |
| 401 | block_size, |
| 402 | ) |
| 403 | |
| 404 | if not self.constant_draft_positions or token_index == 0: |
| 405 | _, per_layer_attn_metadata = ( |
| 406 | self.build_per_group_and_layer_attn_metadata( |
| 407 | common_attn_metadata, draft_index=spec_step_idx |
| 408 | ) |
| 409 | ) |
| 410 | |
| 411 | self.input_ids[:batch_size] = input_ids |
| 412 | self.hidden_states[:batch_size] = hidden_states |
| 413 | if self.supports_mm_inputs: |
| 414 | self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids) |
| 415 | |
| 416 | input_ids = None |
| 417 | inputs_embeds = self.inputs_embeds[:input_batch_size] |
| 418 | else: |
| 419 | input_ids = self.input_ids[:input_batch_size] |
| 420 | inputs_embeds = None |
| 421 | |
| 422 | model_kwargs = { |
| 423 | "input_ids": input_ids, |
| 424 | "positions": self._get_positions(input_batch_size), |
| 425 | "inputs_embeds": inputs_embeds, |
| 426 | "spec_step_idx": spec_step_idx, |
| 427 | } |
| 428 | if self.pass_hidden_states_to_model: |
| 429 | model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] |
| 430 | |
| 431 | with set_forward_context( |
| 432 | per_layer_attn_metadata, |
| 433 | self.vllm_config, |
| 434 | num_tokens=input_batch_size, |
| 435 | num_tokens_across_dp=batch_size_across_dp, |
| 436 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 437 | slot_mapping=self._get_slot_mapping(input_batch_size), |
| 438 | ): |
| 439 | ret_hidden_states = self.model(**model_kwargs) |
| 440 | if not self.model_returns_tuple(): |
| 441 | last_hidden_states = ret_hidden_states |
| 442 | hidden_states = ret_hidden_states |
| 443 | else: |
| 444 | last_hidden_states, hidden_states = ret_hidden_states |
| 445 | |
| 446 | hidden_states = hidden_states[:batch_size] |
| 447 | draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( |
| 448 | last_hidden_states[:batch_size], |
| 449 | sampling_metadata, |
| 450 | spec_step_idx=spec_step_idx, |
| 451 | ) |
| 452 | if draft_probs is not None: |
| 453 | assert draft_probs_list is not None |
| 454 | draft_probs_list.append(draft_probs) |
| 455 | draft_token_ids_list.append(draft_token_ids) |
| 456 | |
| 457 | draft_token_ids = torch.stack(draft_token_ids_list, dim=1) |
| 458 | if draft_probs_list is not None: |
| 459 | self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous() |
| 460 | return draft_token_ids |