| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). |
| 4 | |
| 5 | Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP |
| 6 | depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj`` |
| 7 | (``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the |
| 8 | same short convolutions as the backbone; its attention is full or sliding-window |
| 9 | per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared |
| 10 | ``chain_norm`` is applied after every depth; its output is both the logits input |
| 11 | and the previous hidden state fed to the next depth. |
| 12 | |
| 13 | The draft shares the target's token embedding table and LM head |
| 14 | (``load_eagle_model`` wires those references) and applies the backbone |
| 15 | ``embed_norm`` on top: the depth layers were trained on the same normed |
| 16 | embeddings the backbone consumes (their own ``embed_norm`` weights are |
| 17 | near-identity trims, unlike the backbone's whitening ``embed_norm``). |
| 18 | """ |
| 19 | |
| 20 | from __future__ import annotations |
| 21 | |
| 22 | from collections.abc import Iterable |
| 23 | |
| 24 | import regex as re |
| 25 | import torch |
| 26 | from torch import nn |
| 27 | |
| 28 | from vllm.config import VllmConfig |
| 29 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 30 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 31 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 32 | from vllm.model_executor.model_loader.mtp_validation import ( |
| 33 | is_mtp_completeness_check_enabled, |
| 34 | ) |
| 35 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 36 | from vllm.model_executor.models.interfaces import SupportsMultiModalEmbeddings |
| 37 | from vllm.model_executor.models.utils import maybe_prefix |
| 38 | from vllm.sequence import IntermediateTensors |
| 39 | |
| 40 | from ..configs import InklingModelConfig |
| 41 | from .layernorm import InklingRMSNorm |
| 42 | from .model import InklingDecoderLayer, InklingReplicatedEmbedding |
| 43 | from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm |
| 44 | |
| 45 | # Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. |
| 46 | # Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a |
| 47 | # local (pname, wname, shard) list since the MTP loader remaps by hand. |
| 48 | _ATTENTION_PARAMS_MAPPING = [ |
| 49 | ("qkvr", "wq_du", 0), |
| 50 | ("qkvr", "wk_dv", 1), |
| 51 | ("qkvr", "wv_dv", 2), |
| 52 | ("qkvr", "wr_du", 3), |
| 53 | ] |
| 54 | |
| 55 | |
| 56 | def _mtp_depth_from_name(name: str) -> int | None: |
| 57 | m = re.search(r"\.mtp\.layers\.(\d+)\.", name) |
| 58 | return int(m.group(1)) if m else None |
| 59 | |
| 60 | |
| 61 | def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int: |
| 62 | num_layers = min(n_predict, num_spec) if num_spec else n_predict |
| 63 | if num_layers <= 0: |
| 64 | raise ValueError( |
| 65 | "Inkling MTP requires num_nextn_predict_layers and " |
| 66 | "num_speculative_tokens to select at least one depth layer." |
| 67 | ) |
| 68 | return num_layers |
| 69 | |
| 70 | |
| 71 | class InklingMTPDepthLayer(nn.Module): |
| 72 | """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" |
| 73 | |
| 74 | def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: |
| 75 | super().__init__() |
| 76 | self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 77 | self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 78 | self.input_proj = ReplicatedLinear( |
| 79 | config.hidden_size * 2, |
| 80 | config.hidden_size, |
| 81 | bias=False, |
| 82 | return_bias=False, |
| 83 | prefix=f"{prefix}.input_proj", |
| 84 | ) |
| 85 | # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs |
| 86 | # full attention (the swa_* head config and sliding_window window) to |
| 87 | # match this depth's checkpoint transformer_block weights. |
| 88 | self.transformer_block = InklingDecoderLayer( |
| 89 | config, |
| 90 | layer_id=0, |
| 91 | is_local=is_local, |
| 92 | quant_config=None, |
| 93 | prefix=f"{prefix}.transformer_block", |
| 94 | force_dense_mlp=True, |
| 95 | ) |
| 96 | |
| 97 | def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: |
| 98 | # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)] |
| 99 | # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch. |
| 100 | hidden = self.input_proj(combined) |
| 101 | # The short conv self-fetches its paged SWA-cache metadata from the |
| 102 | # forward context (via its conv_owner prefix); no conv_meta to thread. |
| 103 | return self.transformer_block(positions, hidden) |
| 104 | |
| 105 | |
| 106 | class InklingMultiTokenPredictor(nn.Module): |
| 107 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 108 | super().__init__() |
| 109 | assert vllm_config.speculative_config is not None |
| 110 | config: InklingModelConfig = ( |
| 111 | vllm_config.speculative_config.draft_model_config.hf_config |
| 112 | ) |
| 113 | self.config = config |
| 114 | # The checkpoint ships num_nextn_predict_layers depth blocks, but only |
| 115 | # the first ``num_speculative_tokens`` are exercised (step i uses depth |
| 116 | # i). Build only those to save memory — each depth is a full Inkling block |
| 117 | # with its own (large) full-history sconv caches and KV cache. |
| 118 | n_predict = config.num_nextn_predict_layers |
| 119 | num_spec = vllm_config.speculative_config.num_speculative_tokens |
| 120 | self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec) |
| 121 | self.chain_hidden_post_norm = config.chain_hidden_post_norm |
| 122 | |
| 123 | # Depth blocks whose attention is sliding-window (swa_* head config) |
| 124 | # rather than full; keyed by MTP depth via the checkpoint's |
| 125 | # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors |
| 126 | # InklingModel's local_ids split, but over MTP depths, not backbone |
| 127 | # layers. |
| 128 | local_ids = set(config.local_layer_ids) |
| 129 | |
| 130 | # Keyed by depth index (str) to mirror the checkpoint layout. |
| 131 | self.layers = nn.ModuleDict( |
| 132 | { |
| 133 | str(idx): InklingMTPDepthLayer( |
| 134 | config, f"{prefix}.layers.{idx}", idx in local_ids |
| 135 | ) |
| 136 | for idx in range(self.num_mtp_layers) |
| 137 | } |
| 138 | ) |
| 139 | self.chain_norm = ( |
| 140 | InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 141 | if self.chain_hidden_post_norm |
| 142 | else None |
| 143 | ) |
| 144 | # The target's raw token embedding (pre embed_norm), attached by |
| 145 | # load_eagle_model. Never materialized here: building our own |
| 146 | # replicated copy would transiently double the 2.3 GiB table. |
| 147 | self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment] |
| 148 | # The depth layers consume the *backbone-normed* embedding |
| 149 | # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights |
| 150 | # are near-identity (trained on already-normalized inputs), and |
| 151 | # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70. |
| 152 | # Weight loaded from the target's embed_norm.weight; gated like the |
| 153 | # target's InklingModel.embed_norm. |
| 154 | self.backbone_embed_norm = ( |
| 155 | InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 156 | if config.use_embed_norm |
| 157 | else None |
| 158 | ) |
| 159 | |
| 160 | def embed_input_ids( |
| 161 | self, |
| 162 | input_ids: torch.Tensor, |
| 163 | multimodal_embeddings: object | None = None, |
| 164 | *, |
| 165 | is_multimodal: torch.Tensor | None = None, |
| 166 | ) -> torch.Tensor: |
| 167 | """Draft-prefill embedding: fused gather + backbone embed_norm, then |
| 168 | the target's tower embeddings scattered in unnormed (the backbone |
| 169 | convention — MM embeds are merged after embed_norm).""" |
| 170 | norm = self.backbone_embed_norm |
| 171 | embeds = embed_rmsnorm( |
| 172 | input_ids, |
| 173 | self.embed_tokens.weight, |
| 174 | norm.weight if norm is not None else None, |
| 175 | norm.variance_epsilon if norm is not None else 0.0, |
| 176 | ) |
| 177 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type] |
| 178 | return embeds |
| 179 | from vllm.model_executor.models.utils import _merge_multimodal_embeddings |
| 180 | |
| 181 | assert is_multimodal is not None |
| 182 | return _merge_multimodal_embeddings( |
| 183 | inputs_embeds=embeds, |
| 184 | multimodal_embeddings=multimodal_embeddings, |
| 185 | is_multimodal=is_multimodal, |
| 186 | ) |
| 187 | |
| 188 | def fused_input_cat( |
| 189 | self, |
| 190 | layer: InklingMTPDepthLayer, |
| 191 | previous_hidden: torch.Tensor, |
| 192 | input_ids: torch.Tensor, |
| 193 | inputs_embeds: torch.Tensor | None, |
| 194 | ) -> torch.Tensor: |
| 195 | """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one |
| 196 | launch: embedding row gather + the backbone embed_norm + the depth |
| 197 | embed_norm chain on one side, hidden_norm on the other, written |
| 198 | straight into the cat buffer.""" |
| 199 | hidden_w = layer.hidden_norm.weight |
| 200 | embed_w = layer.embed_norm.weight |
| 201 | eps = layer.hidden_norm.variance_epsilon |
| 202 | if inputs_embeds is not None: |
| 203 | # Draft prefill with target-merged MM embeddings (already |
| 204 | # backbone-normed via embed_input_ids); only the depth embed_norm |
| 205 | # remains. |
| 206 | return embed_dual_rmsnorm_cat( |
| 207 | previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds |
| 208 | ) |
| 209 | return embed_dual_rmsnorm_cat( |
| 210 | previous_hidden, |
| 211 | hidden_w, |
| 212 | embed_w, |
| 213 | eps, |
| 214 | input_ids=input_ids, |
| 215 | embed_table=self.embed_tokens.weight, |
| 216 | pre_norm_weight=( |
| 217 | self.backbone_embed_norm.weight |
| 218 | if self.backbone_embed_norm is not None |
| 219 | else None |
| 220 | ), |
| 221 | ) |
| 222 | |
| 223 | def forward( |
| 224 | self, |
| 225 | input_ids: torch.Tensor, |
| 226 | positions: torch.Tensor, |
| 227 | previous_hidden_states: torch.Tensor, |
| 228 | inputs_embeds: torch.Tensor | None = None, |
| 229 | spec_step_idx: int = 0, |
| 230 | ) -> torch.Tensor: |
| 231 | # The draft's short conv is a paged SWA-cache layer (its conv_owner is |
| 232 | # auto-enumerated as a draft attention layer); its per-token metadata is |
| 233 | # built by the speculator's build_attn_metadata and read from the |
| 234 | # forward context, so nothing extra is threaded here. |
| 235 | depth = spec_step_idx % self.num_mtp_layers |
| 236 | layer = self.layers[str(depth)] |
| 237 | combined = self.fused_input_cat( |
| 238 | layer, previous_hidden_states, input_ids, inputs_embeds |
| 239 | ) |
| 240 | hidden = layer(combined, positions) |
| 241 | if self.chain_norm is not None: |
| 242 | hidden = self.chain_norm(hidden) |
| 243 | return hidden |
| 244 | |
| 245 | |
| 246 | class InklingMTP(nn.Module, SupportsMultiModalEmbeddings): |
| 247 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 248 | super().__init__() |
| 249 | assert vllm_config.speculative_config is not None |
| 250 | config: InklingModelConfig = ( |
| 251 | vllm_config.speculative_config.draft_model_config.hf_config |
| 252 | ) |
| 253 | self.config = config |
| 254 | self.model = InklingMultiTokenPredictor( |
| 255 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 256 | ) |
| 257 | # The target's (vocab-sharded) LM head, attached by load_eagle_model; |
| 258 | # never materialized here (same reasoning as model.embed_tokens). |
| 259 | self.lm_head: ParallelLMHead = None # type: ignore[assignment] |
| 260 | self.logits_processor = LogitsProcessor( |
| 261 | config.padded_vocab_size, |
| 262 | org_vocab_size=config.vocab_size, |
| 263 | soft_cap=config.final_logit_softcapping, |
| 264 | ) |
| 265 | self._logits_zero: torch.Tensor | None = None |
| 266 | |
| 267 | def embed_input_ids( |
| 268 | self, |
| 269 | input_ids: torch.Tensor, |
| 270 | multimodal_embeddings: object | None = None, |
| 271 | *, |
| 272 | is_multimodal: torch.Tensor | None = None, |
| 273 | ) -> torch.Tensor: |
| 274 | return self.model.embed_input_ids( |
| 275 | input_ids, multimodal_embeddings, is_multimodal=is_multimodal |
| 276 | ) |
| 277 | |
| 278 | def forward( |
| 279 | self, |
| 280 | input_ids: torch.Tensor, |
| 281 | positions: torch.Tensor, |
| 282 | hidden_states: torch.Tensor, |
| 283 | intermediate_tensors: IntermediateTensors | None = None, |
| 284 | inputs_embeds: torch.Tensor | None = None, |
| 285 | spec_step_idx: int = 0, |
| 286 | ) -> torch.Tensor: |
| 287 | return self.model( |
| 288 | input_ids, |
| 289 | positions, |
| 290 | hidden_states, |
| 291 | inputs_embeds, |
| 292 | spec_step_idx, |
| 293 | ) |
| 294 | |
| 295 | def compute_logits( |
| 296 | self, |
| 297 | hidden_states: torch.Tensor, |
| 298 | spec_step_idx: int = 0, |
| 299 | ) -> torch.Tensor | None: |
| 300 | # The MTP shares the base model's LM head, which is trained on |
| 301 | # ``hidden / mup``-scaled inputs, so apply the same mup scaling here |
| 302 | # for a matching logit scale — folded into the lm_head GEMM alpha |
| 303 | # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant |
| 304 | # for greedy draft sampling, but it matters for the gumbel sampling |
| 305 | # distribution at temperature > 0.) |
| 306 | mup = self.config.logits_mup_width_multiplier |
| 307 | if not mup: |
| 308 | return self.logits_processor(self.lm_head, hidden_states) |
| 309 | assert self.logits_processor.soft_cap is None |
| 310 | assert self.logits_processor.scale == 1.0 |
| 311 | w = self.lm_head.weight |
| 312 | if self._logits_zero is None: |
| 313 | self._logits_zero = w.new_zeros(1) |
| 314 | logits = torch.addmm( |
| 315 | self._logits_zero, |
| 316 | hidden_states, |
| 317 | w.t(), |
| 318 | beta=0.0, |
| 319 | alpha=1.0 / mup, |
| 320 | ) |
| 321 | logits = self.logits_processor._gather_logits(logits) |
| 322 | if logits is not None: |
| 323 | logits = logits[..., : self.logits_processor.org_vocab_size] |
| 324 | return logits |
| 325 | |
| 326 | def get_top_tokens( |
| 327 | self, |
| 328 | hidden_states: torch.Tensor, |
| 329 | spec_step_idx: int = 0, |
| 330 | ) -> torch.Tensor: |
| 331 | """Greedy draft tokens via rank-local argmax + tiny (value, index) |
| 332 | reduction — no full-vocab logits all-gather. The muP divisor is a |
| 333 | positive scalar, so the argmax is invariant and the scaling is |
| 334 | skipped entirely.""" |
| 335 | return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) |
| 336 | |
| 337 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 338 | return _load_inkling_mtp_weights(self, weights) |
| 339 | |
| 340 | |
| 341 | def _load_inkling_mtp_weights( |
| 342 | module: InklingMTP, |
| 343 | weights: Iterable[tuple[str, torch.Tensor]], |
| 344 | ) -> set[str]: |
| 345 | """Load ``model.mtp.*`` weights into the MTP module. |
| 346 | |
| 347 | Checkpoint keys look like ``model.mtp.chain_norm.weight`` and |
| 348 | ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone |
| 349 | layer's fused-projection layout, so we apply the same qkvr / gate_up / down |
| 350 | remapping as ``_load_inkling_weights``. Token embedding and LM head are shared |
| 351 | (provided by ``load_eagle_model``) and are not present in mtp.safetensors. |
| 352 | """ |
| 353 | # Per-depth attention is full or sliding-window (config.local_layer_ids); |
| 354 | # each depth's qkvr MergedColumnParallelLinear is built with the matching |
| 355 | # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding. |
| 356 | # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never |
| 357 | # exceeds a layer's kv-head count and no GQA K/V replication is needed here. |
| 358 | params = dict(module.named_parameters()) |
| 359 | loaded: set[str] = set() |
| 360 | |
| 361 | def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: |
| 362 | param = params.get(name) |
| 363 | if param is None: |
| 364 | return False |
| 365 | loader = getattr(param, "weight_loader", default_weight_loader) |
| 366 | if shard_id is None: |
| 367 | if loader is default_weight_loader or param.shape == weight.shape: |
| 368 | default_weight_loader(param, weight) |
| 369 | else: |
| 370 | loader(param, weight) |
| 371 | else: |
| 372 | loader(param, weight, shard_id) # type: ignore[call-arg] |
| 373 | loaded.add(name) |
| 374 | return True |
| 375 | |
| 376 | for name, weight in weights: |
| 377 | depth = _mtp_depth_from_name(name) |
| 378 | # Token embedding and LM head are never materialized on the draft |
| 379 | # (no params to load into); load_eagle_model attaches the target's. |
| 380 | if name in ("model.llm.embed.weight", "model.llm.unembed.weight"): |
| 381 | continue |
| 382 | # The backbone embed_norm, applied to the shared embedding before the |
| 383 | # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The |
| 384 | # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded |
| 385 | # below. Only the shared backbone key routes here. |
| 386 | if name == "model.llm.embed_norm.weight": |
| 387 | _load("model.backbone_embed_norm.weight", weight) |
| 388 | continue |
| 389 | # Only consume the MTP weights; everything else belongs to the target. |
| 390 | if ".mtp." not in name: |
| 391 | continue |
| 392 | # Skip depth blocks beyond the ones we built (num_speculative_tokens). |
| 393 | if depth is not None and depth >= module.model.num_mtp_layers: |
| 394 | continue |
| 395 | # model.mtp.chain_norm.weight -> model.chain_norm.weight |
| 396 | # model.mtp.layers.{i}.X -> model.layers.{i}.X |
| 397 | original_name = name |
| 398 | name = name.replace(".mtp.layers.", ".layers.").replace( |
| 399 | ".mtp.chain_norm.", ".chain_norm." |
| 400 | ) |
| 401 | |
| 402 | if ".chain_norm." in name and module.model.chain_norm is None: |
| 403 | raise ValueError( |
| 404 | "Inkling checkpoint contains chain_norm weights but " |
| 405 | "chain_hidden_post_norm is disabled." |
| 406 | ) |
| 407 | |
| 408 | # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr). |
| 409 | matched = False |
| 410 | for pname, wname, shard in _ATTENTION_PARAMS_MAPPING: |
| 411 | if f".attn.{wname}." in name: |
| 412 | mapped_name = name.replace(f".{wname}.", f".{pname}.") |
| 413 | if not _load(mapped_name, weight, shard): |
| 414 | raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") |
| 415 | matched = True |
| 416 | break |
| 417 | if matched: |
| 418 | continue |
| 419 | |
| 420 | # Dense MLP fused gate/up + down. |
| 421 | if ".mlp.w13_dn.weight" in name: |
| 422 | loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) |
| 423 | elif ".mlp.w2_md.weight" in name: |
| 424 | loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) |
| 425 | else: |
| 426 | if name.endswith(".bias") and name not in params: |
| 427 | continue |
| 428 | loaded_weight = _load(name, weight) |
| 429 | if not loaded_weight: |
| 430 | raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") |
| 431 | required = { |
| 432 | name |
| 433 | for name in params |
| 434 | if name.startswith("model.layers.") or name.startswith("model.chain_norm.") |
| 435 | } |
| 436 | if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled(): |
| 437 | raise ValueError( |
| 438 | "Inkling MTP checkpoint is missing required parameters: " |
| 439 | + ", ".join(missing) |
| 440 | ) |
| 441 | return loaded |
| 442 | |
| 443 | |
| 444 | EntryClass = [InklingMTP] |