| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DFlash speculator for Laguna target models. |
| 4 | |
| 5 | Laguna DFlash uses a uniform drafter layer flavor (`layer_types` all full |
| 6 | or all sliding). The draft checkpoint shares token embedding and lm_head |
| 7 | weights with the target model through the generic spec-decode proposer. |
| 8 | """ |
| 9 | |
| 10 | from collections.abc import Iterable |
| 11 | |
| 12 | import torch |
| 13 | from torch import nn |
| 14 | |
| 15 | from vllm import _custom_ops as ops |
| 16 | from vllm.compilation.decorators import support_torch_compile |
| 17 | from vllm.config import VllmConfig |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 20 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 21 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 22 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 23 | ParallelLMHead, |
| 24 | VocabParallelEmbedding, |
| 25 | ) |
| 26 | from vllm.model_executor.model_loader.weight_utils import ( |
| 27 | default_weight_loader, |
| 28 | maybe_remap_kv_scale_name, |
| 29 | ) |
| 30 | from vllm.model_executor.models.interfaces import EagleModelMixin, SupportsEagle3 |
| 31 | from vllm.multimodal.inputs import NestedTensors |
| 32 | |
| 33 | from .laguna import LagunaDecoderLayer |
| 34 | from .qwen3_dflash import DFlashQwen3Model |
| 35 | from .utils import ( |
| 36 | AutoWeightsLoader, |
| 37 | get_draft_quant_config, |
| 38 | maybe_prefix, |
| 39 | process_eagle_weight, |
| 40 | ) |
| 41 | |
| 42 | logger = init_logger(__name__) |
| 43 | |
| 44 | |
| 45 | def _get_dflash_layer_types(config) -> tuple[str, ...]: |
| 46 | layer_types = getattr(config, "layer_types", None) |
| 47 | if layer_types is None: |
| 48 | raise ValueError("Laguna DFlash config requires `layer_types`.") |
| 49 | if len(layer_types) != config.num_hidden_layers: |
| 50 | raise ValueError( |
| 51 | f"DFlash layer_types length {len(layer_types)} does not match " |
| 52 | f"num_hidden_layers {config.num_hidden_layers}." |
| 53 | ) |
| 54 | # Laguna DFlash checkpoints use a uniform drafter attention flavor. |
| 55 | if len(set(layer_types)) > 1: |
| 56 | raise NotImplementedError( |
| 57 | "Laguna DFlash drafter requires a uniform `layer_types` " |
| 58 | f"(got {sorted(set(layer_types))})." |
| 59 | ) |
| 60 | return tuple(layer_types) |
| 61 | |
| 62 | |
| 63 | @support_torch_compile |
| 64 | class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin): |
| 65 | def __init__( |
| 66 | self, |
| 67 | *, |
| 68 | vllm_config: VllmConfig, |
| 69 | prefix: str = "", |
| 70 | ) -> None: |
| 71 | nn.Module.__init__(self) |
| 72 | speculative_config = vllm_config.speculative_config |
| 73 | assert speculative_config is not None |
| 74 | self.config = speculative_config.draft_model_config.hf_config |
| 75 | self.vocab_size = self.config.vocab_size |
| 76 | self.quant_config = get_draft_quant_config(vllm_config) |
| 77 | |
| 78 | target_layer_ids = self.config.dflash_config["target_layer_ids"] |
| 79 | if not target_layer_ids: |
| 80 | raise ValueError( |
| 81 | "Laguna DFlash config requires non-empty " |
| 82 | "`dflash_config.target_layer_ids`." |
| 83 | ) |
| 84 | |
| 85 | self.embed_tokens = VocabParallelEmbedding( |
| 86 | self.config.vocab_size, |
| 87 | self.config.hidden_size, |
| 88 | prefix=maybe_prefix(prefix, "embed_tokens"), |
| 89 | ) |
| 90 | |
| 91 | self.mask_token_id = self.config.dflash_config.get("mask_token_id") |
| 92 | self.register_buffer( |
| 93 | "mask_embedding", |
| 94 | torch.zeros( |
| 95 | self.config.hidden_size, |
| 96 | dtype=vllm_config.model_config.dtype, |
| 97 | ), |
| 98 | persistent=False, |
| 99 | ) |
| 100 | self.has_separate_mask_embedding = False |
| 101 | |
| 102 | self.layer_types = _get_dflash_layer_types(self.config) |
| 103 | target_layer_count = self.config.target_layer_count |
| 104 | self.layers = nn.ModuleList( |
| 105 | [ |
| 106 | LagunaDecoderLayer( |
| 107 | prefix=maybe_prefix(prefix, f"layers.{layer_idx}"), |
| 108 | config=self.config, |
| 109 | cache_config=vllm_config.cache_config, |
| 110 | quant_config=self.quant_config, |
| 111 | layer_idx=layer_idx, |
| 112 | attention_prefix=maybe_prefix( |
| 113 | prefix, f"layers.{layer_idx + target_layer_count}" |
| 114 | ), |
| 115 | ) |
| 116 | for layer_idx in range(self.config.num_hidden_layers) |
| 117 | ] |
| 118 | ) |
| 119 | for layer in self.layers: |
| 120 | if getattr(layer.self_attn, "sliding_window", None) is not None: |
| 121 | # DFlash inserts verifier-context K/V at absolute cache slots. |
| 122 | # Keep full KV allocation; SWA remains a compute-time limit. |
| 123 | layer.self_attn.attn.sliding_window = None |
| 124 | num_features_to_use = len(target_layer_ids) |
| 125 | target_hidden_size = vllm_config.model_config.get_hidden_size() |
| 126 | fc_input_size = target_hidden_size * num_features_to_use |
| 127 | self.num_aux_slices = num_features_to_use |
| 128 | self.aux_hidden_norms = nn.ModuleList( |
| 129 | [ |
| 130 | RMSNorm( |
| 131 | fc_input_size // num_features_to_use, |
| 132 | eps=self.config.rms_norm_eps, |
| 133 | ) |
| 134 | for _ in range(num_features_to_use) |
| 135 | ] |
| 136 | ) |
| 137 | self.fc = ReplicatedLinear( |
| 138 | input_size=fc_input_size, |
| 139 | output_size=self.config.hidden_size, |
| 140 | bias=False, |
| 141 | params_dtype=vllm_config.model_config.dtype, |
| 142 | quant_config=self.quant_config, |
| 143 | prefix=maybe_prefix(prefix, "fc"), |
| 144 | return_bias=False, |
| 145 | ) |
| 146 | self.hidden_norm = RMSNorm( |
| 147 | self.config.hidden_size, |
| 148 | eps=self.config.rms_norm_eps, |
| 149 | ) |
| 150 | self.norm = RMSNorm( |
| 151 | self.config.hidden_size, |
| 152 | eps=self.config.rms_norm_eps, |
| 153 | ) |
| 154 | |
| 155 | def _build_context_kv_buffers( |
| 156 | self, |
| 157 | layers_attn: list[nn.Module], |
| 158 | has_bias: bool, |
| 159 | ) -> None: |
| 160 | self._kv_weights = torch.stack( |
| 161 | [a.qkv_proj.weight[a.q_size :] for a in layers_attn], dim=0 |
| 162 | ).contiguous() |
| 163 | if has_bias: |
| 164 | self._kv_biases: torch.Tensor | None = torch.stack( |
| 165 | [a.qkv_proj.bias[a.q_size :] for a in layers_attn], dim=0 |
| 166 | ).contiguous() |
| 167 | else: |
| 168 | self._kv_biases = None |
| 169 | self._input_layernorm_weights = torch.stack( |
| 170 | [layer.input_layernorm.weight.data for layer in self.layers], dim=0 |
| 171 | ).contiguous() |
| 172 | self._k_norm_weights = torch.stack( |
| 173 | [a.k_norm.weight.data for a in layers_attn], dim=0 |
| 174 | ).contiguous() |
| 175 | |
| 176 | def _project_context_kv( |
| 177 | self, |
| 178 | context_states: torch.Tensor, |
| 179 | num_ctx: int, |
| 180 | num_layers: int, |
| 181 | num_kv_heads: int, |
| 182 | head_dim: int, |
| 183 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 184 | normed_context_states = torch.empty( |
| 185 | (num_layers, num_ctx, context_states.shape[-1]), |
| 186 | dtype=context_states.dtype, |
| 187 | device=context_states.device, |
| 188 | ) |
| 189 | ops.rms_norm( |
| 190 | normed_context_states, |
| 191 | context_states.unsqueeze(0).expand(num_layers, -1, -1), |
| 192 | self._input_layernorm_weights, |
| 193 | self._rms_norm_eps, |
| 194 | ) |
| 195 | all_kv_flat = torch.bmm( |
| 196 | normed_context_states, |
| 197 | self._kv_weights.transpose(1, 2), |
| 198 | ) |
| 199 | if self._kv_biases is not None: |
| 200 | all_kv_flat += self._kv_biases[:, None, :] |
| 201 | all_kv = ( |
| 202 | all_kv_flat.view(num_layers, num_ctx, 2, num_kv_heads, head_dim) |
| 203 | .permute(2, 0, 1, 3, 4) |
| 204 | .contiguous() |
| 205 | ) |
| 206 | all_k = all_kv[0] |
| 207 | all_v = all_kv[1] |
| 208 | return all_k, all_v |
| 209 | |
| 210 | def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor: |
| 211 | all_k_normed = torch.empty_like(all_k) |
| 212 | ops.rms_norm( |
| 213 | all_k_normed, |
| 214 | all_k, |
| 215 | self._k_norm_weights, |
| 216 | self._rms_norm_eps, |
| 217 | ) |
| 218 | return all_k_normed |
| 219 | |
| 220 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 221 | if self.quant_config is not None: |
| 222 | weights = self.quant_config.get_cache_scale_mapper().apply(weights) |
| 223 | params_dict = dict(self.named_parameters()) |
| 224 | loaded_params: set[str] = set() |
| 225 | for name, loaded_weight in weights: |
| 226 | if "scale" in name: |
| 227 | remapped_name = maybe_remap_kv_scale_name(name, params_dict) |
| 228 | if remapped_name is None: |
| 229 | continue |
| 230 | name = remapped_name |
| 231 | param = params_dict[name] |
| 232 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 233 | weight_loader(param, loaded_weight) |
| 234 | loaded_params.add(name) |
| 235 | return loaded_params |
| 236 | |
| 237 | |
| 238 | class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3): |
| 239 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 240 | nn.Module.__init__(self) |
| 241 | speculative_config = vllm_config.speculative_config |
| 242 | assert speculative_config is not None |
| 243 | self.config = speculative_config.draft_model_config.hf_config |
| 244 | if getattr(self.config, "draft_vocab_size", None) is None: |
| 245 | raise ValueError("Laguna DFlash config requires `draft_vocab_size`.") |
| 246 | self.has_own_embed_tokens = False |
| 247 | self.has_own_lm_head = False |
| 248 | target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() |
| 249 | self.config.target_layer_count = target_layer_num |
| 250 | target_vocab_size = vllm_config.model_config.get_vocab_size() |
| 251 | if self.config.draft_vocab_size != target_vocab_size: |
| 252 | raise ValueError( |
| 253 | "Laguna DFlash shares the target lm_head and requires " |
| 254 | "`draft_vocab_size` to match the target vocabulary size " |
| 255 | f"({self.config.draft_vocab_size} != {target_vocab_size})." |
| 256 | ) |
| 257 | self.model = DFlashLagunaModel( |
| 258 | vllm_config=vllm_config, |
| 259 | prefix="model", |
| 260 | ) |
| 261 | |
| 262 | self.lm_head = ParallelLMHead( |
| 263 | self.config.draft_vocab_size, |
| 264 | self.config.hidden_size, |
| 265 | prefix=maybe_prefix(prefix, "lm_head"), |
| 266 | ) |
| 267 | self.logits_processor = LogitsProcessor(self.config.draft_vocab_size) |
| 268 | |
| 269 | def embed_input_ids( |
| 270 | self, |
| 271 | input_ids: torch.Tensor, |
| 272 | multimodal_embeddings: NestedTensors | None = None, |
| 273 | is_multimodal: torch.Tensor | None = None, |
| 274 | ) -> torch.Tensor: |
| 275 | return self.model.embed_input_ids(input_ids) |
| 276 | |
| 277 | def forward( |
| 278 | self, |
| 279 | input_ids: torch.Tensor, |
| 280 | positions: torch.Tensor, |
| 281 | inputs_embeds: torch.Tensor | None = None, |
| 282 | ) -> torch.Tensor: |
| 283 | return self.model(input_ids, positions, inputs_embeds) |
| 284 | |
| 285 | def compute_logits( |
| 286 | self, |
| 287 | hidden_states: torch.Tensor, |
| 288 | ) -> torch.Tensor | None: |
| 289 | return self.logits_processor(self.lm_head, hidden_states) |
| 290 | |
| 291 | def precompute_and_store_context_kv( |
| 292 | self, |
| 293 | context_states: torch.Tensor, |
| 294 | context_positions: torch.Tensor, |
| 295 | context_slot_mapping: torch.Tensor | None = None, |
| 296 | ) -> None: |
| 297 | self.model.precompute_and_store_context_kv( |
| 298 | context_states, context_positions, context_slot_mapping |
| 299 | ) |
| 300 | |
| 301 | def combine_hidden_states( |
| 302 | self, |
| 303 | hidden_states: torch.Tensor, |
| 304 | ) -> torch.Tensor: |
| 305 | # Normalize each verifier hidden-state slice, concatenate them, then |
| 306 | # project into the drafter hidden size used as DFlash context. |
| 307 | needs_squeeze = hidden_states.dim() == 1 |
| 308 | if needs_squeeze: |
| 309 | hidden_states = hidden_states.unsqueeze(0) |
| 310 | num_slices = self.model.num_aux_slices |
| 311 | slice_size = hidden_states.shape[-1] // num_slices |
| 312 | slices = hidden_states.view(hidden_states.shape[0], num_slices, slice_size) |
| 313 | normed = torch.empty_like(slices) |
| 314 | for i, norm in enumerate(self.model.aux_hidden_norms): |
| 315 | normed[:, i, :] = norm(slices[:, i, :]) |
| 316 | hidden_states = normed.reshape(hidden_states.shape[0], -1) |
| 317 | result = self.model.fc(hidden_states) |
| 318 | result = self.model.hidden_norm(result) |
| 319 | if needs_squeeze: |
| 320 | result = result.squeeze(0) |
| 321 | return result |
| 322 | |
| 323 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 324 | model_weights = {} |
| 325 | for name, loaded_weight in weights: |
| 326 | if "lm_head" not in name: |
| 327 | name = "model." + name |
| 328 | model_weights[name] = loaded_weight |
| 329 | process_eagle_weight(self, name) |
| 330 | |
| 331 | loader = AutoWeightsLoader(self) |
| 332 | loaded_weight_names = loader.load_weights(model_weights.items()) |
| 333 | loaded_weight_names.add("lm_head.weight") |
| 334 | loaded_weight_names.add("model.embed_tokens.weight") |
| 335 | self.model._build_fused_kv_buffers() |
| 336 | return loaded_weight_names |