| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import cast |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | |
| 8 | from vllm.config import SpeculativeConfig |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.model_executor.models.interfaces import SupportsEagle3, supports_eagle3 |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | |
| 15 | def set_eagle3_aux_hidden_state_layers( |
| 16 | model: nn.Module, |
| 17 | spec_config: SpeculativeConfig, |
| 18 | ) -> None: |
| 19 | if not supports_eagle3(model): |
| 20 | raise RuntimeError("Model does not support EAGLE3 interface") |
| 21 | if isinstance(model, type): |
| 22 | raise RuntimeError("Expected model instance for EAGLE3 configuration") |
| 23 | eagle3_model = cast(SupportsEagle3, model) |
| 24 | |
| 25 | aux_layers = get_eagle3_aux_layers_from_config(spec_config) |
| 26 | if aux_layers: |
| 27 | logger.info("Using Eagle3 auxiliary layers from config: %s", aux_layers) |
| 28 | else: |
| 29 | aux_layers = eagle3_model.get_eagle3_default_aux_hidden_state_layers() |
| 30 | logger.info("Using Eagle3 auxiliary layers from model: %s", aux_layers) |
| 31 | eagle3_model.set_aux_hidden_state_layers(aux_layers) |
| 32 | reserve_aux_intermediate_tensor_slots(model) |
| 33 | |
| 34 | |
| 35 | def _inner_decoder(model: nn.Module) -> nn.Module | None: |
| 36 | parent_ref = model |
| 37 | if hasattr(model, "get_language_model"): |
| 38 | parent_ref = model.get_language_model() |
| 39 | elif hasattr(model, "language_model"): |
| 40 | parent_ref = model.language_model |
| 41 | return getattr(parent_ref, "model", None) |
| 42 | |
| 43 | |
| 44 | def verify_supports_aux_hidden_states_over_pp(model: nn.Module, method: str) -> None: |
| 45 | inner = _inner_decoder(model) |
| 46 | if not getattr(inner, "supports_aux_hidden_states_over_pp", False): |
| 47 | raise ValueError( |
| 48 | f"{type(model).__name__} does not support {method} with " |
| 49 | "pipeline parallelism" |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | def aux_hidden_state_relay_keys(model: nn.Module) -> tuple[str, ...]: |
| 54 | from vllm.distributed.parallel_state import get_pp_group |
| 55 | |
| 56 | pp = get_pp_group() |
| 57 | if pp.world_size < 2 or pp.is_first_rank or pp.is_last_rank: |
| 58 | return () |
| 59 | inner = _inner_decoder(model) |
| 60 | assert inner is not None |
| 61 | return tuple( |
| 62 | f"{inner.AUX_HIDDEN_STATE_KEY}{i}" for i in range(inner._aux_slot_base_cached) |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | def reserve_aux_intermediate_tensor_slots(model: nn.Module) -> None: |
| 67 | from vllm.distributed.parallel_state import ( |
| 68 | get_pp_group, |
| 69 | model_parallel_is_initialized, |
| 70 | ) |
| 71 | |
| 72 | if not model_parallel_is_initialized(): |
| 73 | return |
| 74 | pp = get_pp_group() |
| 75 | if pp.world_size < 2 or pp.is_first_rank: |
| 76 | return |
| 77 | inner = _inner_decoder(model) |
| 78 | if inner is None or not getattr(inner, "supports_aux_hidden_states_over_pp", False): |
| 79 | return |
| 80 | |
| 81 | num_aux_states = inner._aux_slot_base_cached |
| 82 | if num_aux_states == 0: |
| 83 | return |
| 84 | |
| 85 | key = inner.AUX_HIDDEN_STATE_KEY |
| 86 | hidden_size = inner.config.hidden_size |
| 87 | make_empty = model.make_empty_intermediate_tensors |
| 88 | |
| 89 | def make_empty_with_aux(batch_size, dtype, device): |
| 90 | tensors = make_empty(batch_size, dtype, device) |
| 91 | for i in range(num_aux_states): |
| 92 | tensors[f"{key}{i}"] = torch.zeros( |
| 93 | (batch_size, hidden_size), dtype=dtype, device=device |
| 94 | ) |
| 95 | return tensors |
| 96 | |
| 97 | model.make_empty_intermediate_tensors = make_empty_with_aux |
| 98 | |
| 99 | |
| 100 | def get_eagle3_aux_layers_from_config( |
| 101 | spec_config: SpeculativeConfig, |
| 102 | ) -> tuple[int, ...] | None: |
| 103 | if not (spec_config and spec_config.draft_model_config): |
| 104 | return None |
| 105 | hf_config = spec_config.draft_model_config.hf_config |
| 106 | layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) |
| 107 | if not layer_ids: |
| 108 | dflash_config = getattr(hf_config, "dflash_config", None) |
| 109 | if dflash_config and isinstance(dflash_config, dict): |
| 110 | # Add 1 to convert DFlash's aux layer id semantics |
| 111 | layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])] |
| 112 | if not layer_ids: |
| 113 | dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None) |
| 114 | if dspark_layer_ids: |
| 115 | if getattr(hf_config, "model_type", None) == "deepseek_v41": |
| 116 | # v4.1 reads the attention *inputs* of its target layers, and |
| 117 | # the target model captures the entry stream of layer L when |
| 118 | # idx+1 == L, so the ids are used as-is. (v4's ids are in |
| 119 | # capture-after semantics and keep the +1.) |
| 120 | layer_ids = list(dspark_layer_ids) |
| 121 | else: |
| 122 | layer_ids = [i + 1 for i in dspark_layer_ids] |
| 123 | if not layer_ids: |
| 124 | # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics. |
| 125 | target_layer_ids = getattr(hf_config, "target_layer_ids", None) |
| 126 | if target_layer_ids: |
| 127 | layer_ids = [i + 1 for i in target_layer_ids] |
| 128 | if not layer_ids: |
| 129 | for config_name in ("dflash_config", "eagle_config"): |
| 130 | drafter_config = getattr(hf_config, config_name, None) |
| 131 | if drafter_config and isinstance(drafter_config, dict): |
| 132 | layer_ids = drafter_config.get("layer_ids") |
| 133 | if layer_ids: |
| 134 | break |
| 135 | if layer_ids and isinstance(layer_ids, (list, tuple)): |
| 136 | return tuple(layer_ids) |
| 137 | return None |