diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index ddb8752a80d..7e425d68eec 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -85,6 +85,34 @@ def test_arch_mapping_applies_before_callable_override(): assert seen_architectures == ["MiMoMTPModel"] +@pytest.mark.cpu_test +def test_inkling_override_exposes_only_first_mtp_depth(): + text_config = _make_hf_config( + architectures=["InklingForCausalLM"], + model_type="inkling_model", + local_layer_ids=[1, 3], + ) + config = _make_hf_config( + architectures=["InklingForConditionalGeneration"], + model_type="inkling_mm_model", + text_config=text_config, + mtp_config={ + "num_nextn_predict_layers": 8, + "local_layer_ids": [0, 2, 4], + }, + ) + + out = SpeculativeConfig.hf_config_override(config) + + assert out is text_config + assert out.model_type == "inkling_mtp" + assert out.architectures == ["InklingMTPModel"] + assert out.n_predict == 1 + assert out.num_nextn_predict_layers == 8 + assert out.chain_hidden_post_norm is False + assert out.local_layer_ids == [0, 2, 4] + + def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.num_hidden_layers = 1 return hf_config diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py index 6d0ae51c9d1..42313e7d677 100644 --- a/tests/models/inkling/test_contract_validation.py +++ b/tests/models/inkling/test_contract_validation.py @@ -8,6 +8,7 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser from vllm.models.inkling.configs import ( InklingAudioConfig, + InklingModelConfig, InklingVisionConfig, ) from vllm.models.inkling.nvidia.sconv_swa_attn import ( @@ -34,6 +35,10 @@ def test_inkling_raw_2d_audio_is_rejected_as_ambiguous(): parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32)) +def test_inkling_mtp_chain_norm_is_disabled_by_default(): + assert InklingModelConfig().chain_hidden_post_norm is False + + def test_inkling_supports_piecewise_cudagraphs(): support = InklingSconvMetadataBuilder.get_cudagraph_support assert support(None, None) == AttentionCGSupport.UNIFORM_BATCH diff --git a/tests/models/inkling/test_mtp_input_fusion.py b/tests/models/inkling/test_mtp_input_fusion.py new file mode 100644 index 00000000000..7ab33dd45ae --- /dev/null +++ b/tests/models/inkling/test_mtp_input_fusion.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bit-exactness tests for the fused MTP depth-layer input kernel. + +``embed_dual_rmsnorm_cat`` must match the unfused module sequence exactly: +each rmsnorm computes in fp32 and rounds to bf16 at the same points as the +vendored ``rmsnorm`` kernel (including the bf16 round-trip between the +chained backbone embed_norm and the depth embed_norm), and the fused row +gather matches ``F.embedding``. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("requires CUDA", allow_module_level=True) + +from vllm.models.inkling.nvidia.ops.norm import ( + embed_dual_rmsnorm_cat, + embed_rmsnorm, + rmsnorm, +) + +EPS = 1e-6 +VOCAB = 4096 + + +def _ref(hidden, w_h, w_e, emb, w_pre=None): + if w_pre is not None: + emb = rmsnorm(emb, w_pre, EPS) + return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w_h = torch.randn(n, device=dev).to(torch.bfloat16) + w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) + w_pre = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) + emb = table[ids.long()] + + # Fused gather + chained backbone pre-norm (the decode draft-step path). + out = embed_dual_rmsnorm_cat( + hidden, + w_h, + w_e, + EPS, + input_ids=ids, + embed_table=table, + pre_norm_weight=w_pre, + ) + assert out.shape == (t, 2 * n) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb, w_pre)) + + # Precomputed embeds, no pre-norm (draft prefill with target-merged MM + # embeddings, already backbone-normed). + out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) + + # Fused gather, no pre-norm (use_embed_norm=False). + out = embed_dual_rmsnorm_cat( + hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table + ) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + ref_emb = table[ids.long()] + + # Gather + embed_norm (base model / MTP prefill embed path). + out = embed_rmsnorm(ids, table, w, EPS) + assert out.shape == (t, n) + assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) + + # Pure gather (use_embed_norm=False / replicated module forward). + out = embed_rmsnorm(ids, table, None, EPS) + assert torch.equal(out, ref_emb) + + # Chained first-layer attn_norm (the target text-path forward): one launch + # emits both the residual and layer 0's normed attention input. + w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) + res, attn_in = embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain) + ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb + assert torch.equal(res, ref_res) + assert torch.equal(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) + + # Chained without embed_norm (use_embed_norm=False). + res, attn_in = embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain) + assert torch.equal(res, ref_emb) + assert torch.equal(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) diff --git a/tests/models/registry.py b/tests/models/registry.py index d85227b75c4..a05362f42b5 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1656,6 +1656,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "tencent/Hy3-preview", speculative_model="tencent/Hy3-preview", ), + "InklingMTPModel": _HfExamplesInfo( + "thinkingmachines/Inkling-NVFP4", + speculative_model="thinkingmachines/Inkling-NVFP4", + tokenizer_mode="inkling", + trust_remote_code=True, + max_model_len=4096, + ), "LongCatFlashMTPModel": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 0cc09f8295d..cf6f0a5ae6b 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -55,6 +55,7 @@ MTPModelTypes = Literal[ "step3p5_mtp", "hy_v3_mtp", "gemma4_mtp", + "inkling_mtp", ] NgramGPUTypes = Literal["ngram_gpu"] DFlashModelTypes = Literal["dflash"] @@ -553,6 +554,26 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]} ) + if hf_config.model_type in ("inkling_mm_model", "inkling_model"): + mtp_config = getattr(hf_config, "mtp_config", None) or {} + hf_config = getattr(hf_config, "text_config", hf_config) + checkpoint_depths = mtp_config.get("num_nextn_predict_layers", 0) + if checkpoint_depths < 1: + raise ValueError("The Inkling checkpoint does not contain MTP weights") + hf_config.model_type = "inkling_mtp" + hf_config.update( + { + # Inkling currently exposes only the first checkpoint depth. + "n_predict": 1, + "num_nextn_predict_layers": checkpoint_depths, + "chain_hidden_post_norm": mtp_config.get( + "chain_hidden_post_norm", False + ), + "local_layer_ids": mtp_config.get("local_layer_ids", []), + "architectures": ["InklingMTPModel"], + } + ) + if hf_config.model_type in ("gemma4_assistant", "gemma4_unified_assistant"): hf_config.model_type = "gemma4_mtp" text_config = getattr(hf_config, "text_config", hf_config) @@ -874,7 +895,7 @@ class SpeculativeConfig: if ( self.num_speculative_tokens > 1 and self.draft_model_config.hf_config.model_type - != "step3p5_mtp" + not in ("step3p5_mtp", "inkling_mtp") ): logger.warning( "Enabling num_speculative_tokens > 1 will run " @@ -953,6 +974,14 @@ class SpeculativeConfig: "`num_speculative_tokens` was not provided" ) + if ( + self.draft_model_config.hf_config.model_type == "inkling_mtp" + and self.num_speculative_tokens != 1 + ): + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + if self.method == "dspark": # DSpark is a semi-autoregressive *block* drafter. A # speculative length smaller than the checkpoint's block diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index f020abdd980..b5ec2519d07 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -630,6 +630,7 @@ _SPECULATIVE_DECODING_MODELS = { "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"), "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"), + "InklingMTPModel": ("vllm.models.inkling", "InklingMTP"), "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), diff --git a/vllm/models/inkling/__init__.py b/vllm/models/inkling/__init__.py index 21385d8fe73..32e58905e25 100644 --- a/vllm/models/inkling/__init__.py +++ b/vllm/models/inkling/__init__.py @@ -7,14 +7,20 @@ if TYPE_CHECKING: InklingForCausalLM, InklingForConditionalGeneration, ) + from .nvidia.mtp import InklingMTP __all__ = [ "InklingForConditionalGeneration", "InklingForCausalLM", + "InklingMTP", ] def __getattr__(name: str): + if name == "InklingMTP": + from .nvidia import mtp + + return mtp.InklingMTP if name in __all__: from .nvidia import model diff --git a/vllm/models/inkling/configs.py b/vllm/models/inkling/configs.py index f5cbff15cc8..aed70b14302 100644 --- a/vllm/models/inkling/configs.py +++ b/vllm/models/inkling/configs.py @@ -61,6 +61,8 @@ class InklingModelConfig(PretrainedConfig): logits_mup_width_multiplier: float | None = None, final_logit_softcapping: float | None = None, tie_word_embeddings: bool = False, + num_nextn_predict_layers: int = 0, + chain_hidden_post_norm: bool = False, **kwargs: Any, ) -> None: if head_dim is None: @@ -135,6 +137,10 @@ class InklingModelConfig(PretrainedConfig): self.unpadded_vocab_size = self.vocab_size self.logits_mup_width_multiplier = logits_mup_width_multiplier self.final_logit_softcapping = final_logit_softcapping + # MTP (multi-token prediction) draft head: number of depth layers in the + # checkpoint (0 if absent); chain_norm applied after each depth. + self.num_nextn_predict_layers = num_nextn_predict_layers + self.chain_hidden_post_norm = chain_hidden_post_norm super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index 0f62af3a4f3..39f8ea5375d 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -123,6 +123,7 @@ class InklingDecoderLayer(nn.Module): quant_config: QuantizationConfig | None, prefix: str, nvfp4_config: InklingNvfp4Config | None = None, + force_dense_mlp: bool = False, ) -> None: super().__init__() # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv @@ -161,7 +162,7 @@ class InklingDecoderLayer(nn.Module): conv_owner=self.conv_state, ) self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - if layer_id < config.dense_mlp_idx: + if force_dense_mlp or layer_id < config.dense_mlp_idx: self.mlp: nn.Module = InklingDenseMLP( hidden_size=config.hidden_size, intermediate_size=config.dense_intermediate_size, @@ -199,9 +200,12 @@ class InklingDecoderLayer(nn.Module): positions: torch.Tensor, hidden_states: torch.Tensor, pending: tuple[torch.Tensor | None, InklingShortConv] | None = None, + defer_mlp_add: bool = False, attn_in: torch.Tensor | None = None, log_scaling: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]]: + ) -> ( + torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]] + ): # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm. # A None delta means the partials sit in the NVLS symm buffer. @@ -219,9 +223,13 @@ class InklingDecoderLayer(nn.Module): attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions ) mlp_output = self.mlp(mlp_in) - # The caller folds mlp_output (pre-reduce, pre-sconv) into the next - # fused sconv+add+rmsnorm. - return hidden_states, (mlp_output, self.mlp_sconv) + if defer_mlp_add: + # Caller folds mlp_output (pre-reduce, pre-sconv) into the next + # fused sconv+add+rmsnorm. + return hidden_states, (mlp_output, self.mlp_sconv) + return _sconv_add_norm( + mlp_output, hidden_states, self.mlp_sconv, None, positions + )[1] class InklingReplicatedEmbedding(nn.Module): @@ -330,6 +338,7 @@ class InklingModel(nn.Module): positions, hidden_states, pending=pending, + defer_mlp_add=True, attn_in=attn_in0, log_scaling=log_scaling, ) diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py new file mode 100644 index 00000000000..868ede3bed6 --- /dev/null +++ b/vllm/models/inkling/nvidia/mtp.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). + +Implements the first MTP depth from the reference ``mtp_model.py`` shipped with +the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` +input projection, and a full Inkling transformer block with a dense bf16 MLP. + +The draft shares the target's token embedding table and LM head +(``load_eagle_model`` wires those references) and applies the backbone +``embed_norm`` on top: the depth layers were trained on the same normed +embeddings the backbone consumes (their own ``embed_norm`` weights are +near-identity trims, unlike the backbone's whitening ``embed_norm``). +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import maybe_prefix +from vllm.sequence import IntermediateTensors + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm + +# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. +# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a +# local (pname, wname, shard) list since the MTP loader remaps by hand. +_ATTENTION_PARAMS_MAPPING = [ + ("qkvr", "wq_du", 0), + ("qkvr", "wk_dv", 1), + ("qkvr", "wv_dv", 2), + ("qkvr", "wr_du", 3), +] + + +def _mtp_depth_from_name(name: str) -> int | None: + m = re.search(r"\.mtp\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +class InklingMTPDepthLayer(nn.Module): + """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" + + def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + super().__init__() + self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + return_bias=False, + prefix=f"{prefix}.input_proj", + ) + # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # full attention (the swa_* head config and sliding_window window) to + # match this depth's checkpoint transformer_block weights. + self.transformer_block = InklingDecoderLayer( + config, + layer_id=0, + is_local=is_local, + quant_config=None, + prefix=f"{prefix}.transformer_block", + nvfp4_config=None, + force_dense_mlp=True, + ) + + def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)] + # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch. + hidden = self.input_proj(combined) + # The short conv self-fetches its paged SWA-cache metadata from the + # forward context (via its conv_owner prefix); no conv_meta to thread. + return self.transformer_block(positions, hidden) + + +class InklingMultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + if vllm_config.speculative_config.num_speculative_tokens != 1: + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + self.chain_hidden_post_norm = config.chain_hidden_post_norm + local_ids = set(config.local_layer_ids) + self.layers = nn.ModuleDict( + {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + ) + self.chain_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.chain_hidden_post_norm + else None + ) + # The target's raw token embedding (pre embed_norm), attached by + # load_eagle_model. Never materialized here: building our own + # replicated copy would transiently double the 2.3 GiB table. + self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment] + # The depth layers consume the *backbone-normed* embedding + # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights + # are near-identity (trained on already-normalized inputs), and + # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70. + # Weight loaded from the target's embed_norm.weight; gated like the + # target's InklingModel.embed_norm. + self.backbone_embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + """Draft-prefill embedding: fused gather + backbone embed_norm, then + the target's tower embeddings scattered in unnormed (the backbone + convention — MM embeds are merged after embed_norm).""" + norm = self.backbone_embed_norm + embeds = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type] + return embeds + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def fused_input_cat( + self, + layer: InklingMTPDepthLayer, + previous_hidden: torch.Tensor, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one + launch: embedding row gather + the backbone embed_norm + the depth + embed_norm chain on one side, hidden_norm on the other, written + straight into the cat buffer.""" + hidden_w = layer.hidden_norm.weight + embed_w = layer.embed_norm.weight + eps = layer.hidden_norm.variance_epsilon + if inputs_embeds is not None: + # Draft prefill with target-merged MM embeddings (already + # backbone-normed via embed_input_ids); only the depth embed_norm + # remains. + return embed_dual_rmsnorm_cat( + previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds + ) + return embed_dual_rmsnorm_cat( + previous_hidden, + hidden_w, + embed_w, + eps, + input_ids=input_ids, + embed_table=self.embed_tokens.weight, + pre_norm_weight=( + self.backbone_embed_norm.weight + if self.backbone_embed_norm is not None + else None + ), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # The draft's short conv is a paged SWA-cache layer (its conv_owner is + # auto-enumerated as a draft attention layer); its per-token metadata is + # built by the speculator's build_attn_metadata and read from the + # forward context, so nothing extra is threaded here. + if spec_step_idx != 0: + raise ValueError("Inkling MTP only supports spec_step_idx=0") + layer = self.layers["0"] + combined = self.fused_input_cat( + layer, previous_hidden_states, input_ids, inputs_embeds + ) + hidden = layer(combined, positions) + if self.chain_norm is not None: + hidden = self.chain_norm(hidden) + return hidden + + +class InklingMTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + self.model = InklingMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # The target's (vocab-sharded) LM head, attached by load_eagle_model; + # never materialized here (same reasoning as model.embed_tokens). + self.lm_head: ParallelLMHead = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor( + config.padded_vocab_size, + org_vocab_size=config.vocab_size, + soft_cap=config.final_logit_softcapping, + ) + self._logits_zero: torch.Tensor | None = None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.embed_input_ids( + input_ids, multimodal_embeddings, is_multimodal=is_multimodal + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + # The MTP shares the base model's LM head, which is trained on + # ``hidden / mup``-scaled inputs, so apply the same mup scaling here + # for a matching logit scale — folded into the lm_head GEMM alpha + # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant + # for greedy draft sampling, but it matters for the gumbel sampling + # distribution at temperature > 0.) + mup = self.config.logits_mup_width_multiplier + if not mup: + return self.logits_processor(self.lm_head, hidden_states) + assert self.logits_processor.soft_cap is None + assert self.logits_processor.scale == 1.0 + w = self.lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self.logits_processor._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + return logits + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Greedy draft tokens via rank-local argmax + tiny (value, index) + reduction — no full-vocab logits all-gather. The muP divisor is a + positive scalar, so the argmax is invariant and the scaling is + skipped entirely.""" + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_mtp_weights(self, weights) + + +def _load_inkling_mtp_weights( + module: InklingMTP, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Load ``model.mtp.*`` weights into the MTP module. + + Checkpoint keys look like ``model.mtp.chain_norm.weight`` and + ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone + layer's fused-projection layout, so we apply the same qkvr / gate_up / down + remapping as ``_load_inkling_weights``. Token embedding and LM head are shared + (provided by ``load_eagle_model``) and are not present in mtp.safetensors. + """ + # Per-depth attention is full or sliding-window (config.local_layer_ids); + # each depth's qkvr MergedColumnParallelLinear is built with the matching + # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding. + # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never + # exceeds a layer's kv-head count and no GQA K/V replication is needed here. + params = dict(module.named_parameters()) + loaded: set[str] = set() + + def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: + param = params.get(name) + if param is None: + return False + loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + if loader is default_weight_loader or param.shape == weight.shape: + default_weight_loader(param, weight) + else: + loader(param, weight) + else: + loader(param, weight, shard_id) # type: ignore[call-arg] + loaded.add(name) + return True + + for name, weight in weights: + depth = _mtp_depth_from_name(name) + # Token embedding and LM head are never materialized on the draft + # (no params to load into); load_eagle_model attaches the target's. + if name in ("model.llm.embed.weight", "model.llm.unembed.weight"): + continue + # The backbone embed_norm, applied to the shared embedding before the + # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The + # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded + # below. Only the shared backbone key routes here. + if name == "model.llm.embed_norm.weight": + _load("model.backbone_embed_norm.weight", weight) + continue + # Only consume the MTP weights; everything else belongs to the target. + if ".mtp." not in name: + continue + # Only the first checkpoint depth is used for MTP=1. + if depth is not None and depth != 0: + continue + # model.mtp.chain_norm.weight -> model.chain_norm.weight + # model.mtp.layers.{i}.X -> model.layers.{i}.X + original_name = name + name = name.replace(".mtp.layers.", ".layers.").replace( + ".mtp.chain_norm.", ".chain_norm." + ) + + if ".chain_norm." in name and module.model.chain_norm is None: + raise ValueError( + "Inkling checkpoint contains chain_norm weights but " + "chain_hidden_post_norm is disabled." + ) + + # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr). + matched = False + for pname, wname, shard in _ATTENTION_PARAMS_MAPPING: + if f".attn.{wname}." in name: + mapped_name = name.replace(f".{wname}.", f".{pname}.") + if not _load(mapped_name, weight, shard): + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + matched = True + break + if matched: + continue + + # Dense MLP fused gate/up + down. + if ".mlp.w13_dn.weight" in name: + loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) + elif ".mlp.w2_md.weight" in name: + loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) + else: + if name.endswith(".bias") and name not in params: + continue + loaded_weight = _load(name, weight) + if not loaded_weight: + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + required = { + name + for name in params + if name.startswith("model.layers.") or name.startswith("model.chain_norm.") + } + if missing := sorted(required - loaded): + raise ValueError( + "Inkling MTP checkpoint is missing required parameters: " + + ", ".join(missing) + ) + return loaded + + +EntryClass = [InklingMTP] diff --git a/vllm/models/inkling/nvidia/ops/norm.py b/vllm/models/inkling/nvidia/ops/norm.py index d58a39a06bf..d34c1c3259b 100644 --- a/vllm/models/inkling/nvidia/ops/norm.py +++ b/vllm/models/inkling/nvidia/ops/norm.py @@ -268,6 +268,105 @@ def embed_rmsnorm( return out +@triton.jit +def _embed_dual_rmsnorm_cat_kernel( + hidden_ptr, # [T, N] + emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER + ids_ptr, # [T] token ids (GATHER only) + w_hidden_ptr, # [N] + w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only) + w_embed_ptr, # [N] + out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))] + eps, + hidden_stride_0, + emb_stride_0, + n_cols, + block_size_n: tl.constexpr, + GATHER: tl.constexpr, + HAS_PRE_NORM: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + if which == 0: + x = tl.load( + hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0 + ).to(tl.float32) + w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + else: + row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m + x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + if HAS_PRE_NORM: + w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + # Round-trip through the output dtype so the chained norm is + # bit-exact vs the unfused pair (which stores bf16 in between). + x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32) + w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + tl.store( + out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n, + (x * rstd * w).to(out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_dual_rmsnorm_cat( + hidden: torch.Tensor, + hidden_weight: torch.Tensor, + embed_weight: torch.Tensor, + eps: float, + *, + embeds: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + embed_table: torch.Tensor | None = None, + pre_norm_weight: torch.Tensor | None = None, +) -> torch.Tensor: + """The MTP depth-layer input in one launch: + ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``. + + The embed side is either a fused row gather ``embed_table[input_ids]`` + (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged + multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the + backbone embed_norm in front of the depth embed_norm (bit-exact vs the + unfused sequence). The concat copies collapse into direct writes.""" + T, n = hidden.shape + if embeds is not None: + assert embeds.shape == hidden.shape + src, ids, src_stride = embeds, embeds, embeds.stride(0) + gather = False + else: + assert input_ids is not None and embed_table is not None + assert input_ids.shape == (T,) and embed_table.shape[1] == n + src, ids, src_stride = embed_table, input_ids, embed_table.stride(0) + gather = True + out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device) + if T == 0: + return out + block_size_n = triton.next_power_of_2(n) + _embed_dual_rmsnorm_cat_kernel[(T, 2)]( + hidden, + src, + ids, + hidden_weight, + pre_norm_weight if pre_norm_weight is not None else embed_weight, + embed_weight, + out, + eps, + hidden.stride(0), + src_stride, + n, + block_size_n, + GATHER=gather, + HAS_PRE_NORM=pre_norm_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + return out + + def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: assert x.ndim == 2, f"{x.shape=}" assert weight.ndim == 1, f"{weight.shape=}"