[Model] Add Inkling MTP=1 support [3/N] (#48869)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <mozf@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

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

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

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)

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,

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

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"),

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
 

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)
 

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,
             )

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]

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=}"