[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>
Woosuk Kwon committed
Jul 16, 2026 at 13:27 UTC
fb5ec0dc9edfdba54882575e33e2a4215bd295b9
11 files changed
+710
-6
tests/config/test_speculative_draft_hf_overrides.py
+28
@@ -85,6 +85,34 @@ def test_arch_mapping_applies_before_callable_override():
85
assert seen_architectures == ["MiMoMTPModel"]
86
87
88
+@pytest.mark.cpu_test
89
+def test_inkling_override_exposes_only_first_mtp_depth():
90
+ text_config = _make_hf_config(
91
+ architectures=["InklingForCausalLM"],
92
+ model_type="inkling_model",
93
+ local_layer_ids=[1, 3],
94
+ )
95
+ config = _make_hf_config(
96
+ architectures=["InklingForConditionalGeneration"],
97
+ model_type="inkling_mm_model",
98
+ text_config=text_config,
99
+ mtp_config={
100
+ "num_nextn_predict_layers": 8,
101
+ "local_layer_ids": [0, 2, 4],
102
+ },
103
+ )
104
+
105
+ out = SpeculativeConfig.hf_config_override(config)
106
+
107
+ assert out is text_config
108
+ assert out.model_type == "inkling_mtp"
109
+ assert out.architectures == ["InklingMTPModel"]
110
+ assert out.n_predict == 1
111
+ assert out.num_nextn_predict_layers == 8
112
+ assert out.chain_hidden_post_norm is False
113
+ assert out.local_layer_ids == [0, 2, 4]
114
+
115
+
116
def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
117
hf_config.num_hidden_layers = 1
118
return hf_config
tests/models/inkling/test_contract_validation.py
+5
@@ -8,6 +8,7 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode
8
from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser
9
from vllm.models.inkling.configs import (
10
InklingAudioConfig,
11
+ InklingModelConfig,
12
InklingVisionConfig,
13
)
14
from vllm.models.inkling.nvidia.sconv_swa_attn import (
@@ -34,6 +35,10 @@ def test_inkling_raw_2d_audio_is_rejected_as_ambiguous():
35
parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32))
36
37
38
+def test_inkling_mtp_chain_norm_is_disabled_by_default():
39
+ assert InklingModelConfig().chain_hidden_post_norm is False
40
+
41
+
42
def test_inkling_supports_piecewise_cudagraphs():
43
support = InklingSconvMetadataBuilder.get_cudagraph_support
44
assert support(None, None) == AttentionCGSupport.UNIFORM_BATCH
tests/models/inkling/test_mtp_input_fusion.py
new
+106
@@ -0,0 +1,106 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Bit-exactness tests for the fused MTP depth-layer input kernel.
4
+
5
+``embed_dual_rmsnorm_cat`` must match the unfused module sequence exactly:
6
+each rmsnorm computes in fp32 and rounds to bf16 at the same points as the
7
+vendored ``rmsnorm`` kernel (including the bf16 round-trip between the
8
+chained backbone embed_norm and the depth embed_norm), and the fused row
9
+gather matches ``F.embedding``.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm.platforms import current_platform
16
+
17
+if not current_platform.is_cuda():
18
+ pytest.skip("requires CUDA", allow_module_level=True)
19
+
20
+from vllm.models.inkling.nvidia.ops.norm import (
21
+ embed_dual_rmsnorm_cat,
22
+ embed_rmsnorm,
23
+ rmsnorm,
24
+)
25
+
26
+EPS = 1e-6
27
+VOCAB = 4096
28
+
29
+
30
+def _ref(hidden, w_h, w_e, emb, w_pre=None):
31
+ if w_pre is not None:
32
+ emb = rmsnorm(emb, w_pre, EPS)
33
+ return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1)
34
+
35
+
36
+@pytest.mark.parametrize("n", [1536, 6144])
37
+@pytest.mark.parametrize("t", [0, 1, 7, 256])
38
+@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
39
+def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None:
40
+ torch.manual_seed(0)
41
+ dev = "cuda"
42
+ table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
43
+ w_h = torch.randn(n, device=dev).to(torch.bfloat16)
44
+ w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16)
45
+ w_pre = torch.randn(n, device=dev).to(torch.bfloat16)
46
+ ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
47
+ hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16)
48
+ emb = table[ids.long()]
49
+
50
+ # Fused gather + chained backbone pre-norm (the decode draft-step path).
51
+ out = embed_dual_rmsnorm_cat(
52
+ hidden,
53
+ w_h,
54
+ w_e,
55
+ EPS,
56
+ input_ids=ids,
57
+ embed_table=table,
58
+ pre_norm_weight=w_pre,
59
+ )
60
+ assert out.shape == (t, 2 * n)
61
+ assert torch.equal(out, _ref(hidden, w_h, w_e, emb, w_pre))
62
+
63
+ # Precomputed embeds, no pre-norm (draft prefill with target-merged MM
64
+ # embeddings, already backbone-normed).
65
+ out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb)
66
+ assert torch.equal(out, _ref(hidden, w_h, w_e, emb))
67
+
68
+ # Fused gather, no pre-norm (use_embed_norm=False).
69
+ out = embed_dual_rmsnorm_cat(
70
+ hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table
71
+ )
72
+ assert torch.equal(out, _ref(hidden, w_h, w_e, emb))
73
+
74
+
75
+@pytest.mark.parametrize("n", [1536, 6144])
76
+@pytest.mark.parametrize("t", [0, 1, 7, 256])
77
+@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
78
+def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None:
79
+ torch.manual_seed(0)
80
+ dev = "cuda"
81
+ table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
82
+ w = torch.randn(n, device=dev).to(torch.bfloat16)
83
+ ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
84
+ ref_emb = table[ids.long()]
85
+
86
+ # Gather + embed_norm (base model / MTP prefill embed path).
87
+ out = embed_rmsnorm(ids, table, w, EPS)
88
+ assert out.shape == (t, n)
89
+ assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb)
90
+
91
+ # Pure gather (use_embed_norm=False / replicated module forward).
92
+ out = embed_rmsnorm(ids, table, None, EPS)
93
+ assert torch.equal(out, ref_emb)
94
+
95
+ # Chained first-layer attn_norm (the target text-path forward): one launch
96
+ # emits both the residual and layer 0's normed attention input.
97
+ w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16)
98
+ res, attn_in = embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain)
99
+ ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb
100
+ assert torch.equal(res, ref_res)
101
+ assert torch.equal(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res)
102
+
103
+ # Chained without embed_norm (use_embed_norm=False).
104
+ res, attn_in = embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain)
105
+ assert torch.equal(res, ref_emb)
106
+ assert torch.equal(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb)
tests/models/registry.py
+7
@@ -1656,6 +1656,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1656
"tencent/Hy3-preview",
1657
speculative_model="tencent/Hy3-preview",
1658
),
1659
+ "InklingMTPModel": _HfExamplesInfo(
1660
+ "thinkingmachines/Inkling-NVFP4",
1661
+ speculative_model="thinkingmachines/Inkling-NVFP4",
1662
+ tokenizer_mode="inkling",
1663
+ trust_remote_code=True,
1664
+ max_model_len=4096,
1665
+ ),
1666
"LongCatFlashMTPModel": _HfExamplesInfo(
1667
"meituan-longcat/LongCat-Flash-Chat",
1668
trust_remote_code=True,
vllm/config/speculative.py
+30
-1
@@ -55,6 +55,7 @@ MTPModelTypes = Literal[
55
"step3p5_mtp",
56
"hy_v3_mtp",
57
"gemma4_mtp",
58
+ "inkling_mtp",
59
]
60
NgramGPUTypes = Literal["ngram_gpu"]
61
DFlashModelTypes = Literal["dflash"]
@@ -553,6 +554,26 @@ class SpeculativeConfig:
554
{"n_predict": n_predict, "architectures": ["HYV3MTPModel"]}
555
)
556
557
+ if hf_config.model_type in ("inkling_mm_model", "inkling_model"):
558
+ mtp_config = getattr(hf_config, "mtp_config", None) or {}
559
+ hf_config = getattr(hf_config, "text_config", hf_config)
560
+ checkpoint_depths = mtp_config.get("num_nextn_predict_layers", 0)
561
+ if checkpoint_depths < 1:
562
+ raise ValueError("The Inkling checkpoint does not contain MTP weights")
563
+ hf_config.model_type = "inkling_mtp"
564
+ hf_config.update(
565
+ {
566
+ # Inkling currently exposes only the first checkpoint depth.
567
+ "n_predict": 1,
568
+ "num_nextn_predict_layers": checkpoint_depths,
569
+ "chain_hidden_post_norm": mtp_config.get(
570
+ "chain_hidden_post_norm", False
571
+ ),
572
+ "local_layer_ids": mtp_config.get("local_layer_ids", []),
573
+ "architectures": ["InklingMTPModel"],
574
+ }
575
+ )
576
+
577
if hf_config.model_type in ("gemma4_assistant", "gemma4_unified_assistant"):
578
hf_config.model_type = "gemma4_mtp"
579
text_config = getattr(hf_config, "text_config", hf_config)
@@ -874,7 +895,7 @@ class SpeculativeConfig:
895
if (
896
self.num_speculative_tokens > 1
897
and self.draft_model_config.hf_config.model_type
877
- != "step3p5_mtp"
898
+ not in ("step3p5_mtp", "inkling_mtp")
899
):
900
logger.warning(
901
"Enabling num_speculative_tokens > 1 will run "
@@ -953,6 +974,14 @@ class SpeculativeConfig:
974
"`num_speculative_tokens` was not provided"
975
)
976
977
+ if (
978
+ self.draft_model_config.hf_config.model_type == "inkling_mtp"
979
+ and self.num_speculative_tokens != 1
980
+ ):
981
+ raise ValueError(
982
+ "Inkling MTP currently supports exactly one speculative token"
983
+ )
984
+
985
if self.method == "dspark":
986
# DSpark is a semi-autoregressive *block* drafter. A
987
# speculative length smaller than the checkpoint's block
vllm/model_executor/models/registry.py
+1
@@ -630,6 +630,7 @@ _SPECULATIVE_DECODING_MODELS = {
630
"DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
631
"MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
632
"BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"),
633
+ "InklingMTPModel": ("vllm.models.inkling", "InklingMTP"),
634
"Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
635
"ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
636
"ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
vllm/models/inkling/__init__.py
+6
@@ -7,14 +7,20 @@ if TYPE_CHECKING:
7
InklingForCausalLM,
8
InklingForConditionalGeneration,
9
)
10
+ from .nvidia.mtp import InklingMTP
11
12
__all__ = [
13
"InklingForConditionalGeneration",
14
"InklingForCausalLM",
15
+ "InklingMTP",
16
]
17
18
19
def __getattr__(name: str):
20
+ if name == "InklingMTP":
21
+ from .nvidia import mtp
22
+
23
+ return mtp.InklingMTP
24
if name in __all__:
25
from .nvidia import model
26
vllm/models/inkling/configs.py
+6
@@ -61,6 +61,8 @@ class InklingModelConfig(PretrainedConfig):
61
logits_mup_width_multiplier: float | None = None,
62
final_logit_softcapping: float | None = None,
63
tie_word_embeddings: bool = False,
64
+ num_nextn_predict_layers: int = 0,
65
+ chain_hidden_post_norm: bool = False,
66
**kwargs: Any,
67
) -> None:
68
if head_dim is None:
@@ -135,6 +137,10 @@ class InklingModelConfig(PretrainedConfig):
137
self.unpadded_vocab_size = self.vocab_size
138
self.logits_mup_width_multiplier = logits_mup_width_multiplier
139
self.final_logit_softcapping = final_logit_softcapping
140
+ # MTP (multi-token prediction) draft head: number of depth layers in the
141
+ # checkpoint (0 if absent); chain_norm applied after each depth.
142
+ self.num_nextn_predict_layers = num_nextn_predict_layers
143
+ self.chain_hidden_post_norm = chain_hidden_post_norm
144
145
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
146
vllm/models/inkling/nvidia/model.py
+14
-5
@@ -123,6 +123,7 @@ class InklingDecoderLayer(nn.Module):
123
quant_config: QuantizationConfig | None,
124
prefix: str,
125
nvfp4_config: InklingNvfp4Config | None = None,
126
+ force_dense_mlp: bool = False,
127
) -> None:
128
super().__init__()
129
# Per-layer owner of the conv state as a paged SWA cache. The 4 sconv
@@ -161,7 +162,7 @@ class InklingDecoderLayer(nn.Module):
162
conv_owner=self.conv_state,
163
)
164
self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
164
- if layer_id < config.dense_mlp_idx:
165
+ if force_dense_mlp or layer_id < config.dense_mlp_idx:
166
self.mlp: nn.Module = InklingDenseMLP(
167
hidden_size=config.hidden_size,
168
intermediate_size=config.dense_intermediate_size,
@@ -199,9 +200,12 @@ class InklingDecoderLayer(nn.Module):
200
positions: torch.Tensor,
201
hidden_states: torch.Tensor,
202
pending: tuple[torch.Tensor | None, InklingShortConv] | None = None,
203
+ defer_mlp_add: bool = False,
204
attn_in: torch.Tensor | None = None,
205
log_scaling: torch.Tensor | None = None,
204
- ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]]:
206
+ ) -> (
207
+ torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]]
208
+ ):
209
# The previous sublayer's (pre-reduce, pre-sconv) delta is folded in
210
# fused with its RS/sconv/AG and this layer's pre-attention rmsnorm.
211
# A None delta means the partials sit in the NVLS symm buffer.
@@ -219,9 +223,13 @@ class InklingDecoderLayer(nn.Module):
223
attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions
224
)
225
mlp_output = self.mlp(mlp_in)
222
- # The caller folds mlp_output (pre-reduce, pre-sconv) into the next
223
- # fused sconv+add+rmsnorm.
224
- return hidden_states, (mlp_output, self.mlp_sconv)
226
+ if defer_mlp_add:
227
+ # Caller folds mlp_output (pre-reduce, pre-sconv) into the next
228
+ # fused sconv+add+rmsnorm.
229
+ return hidden_states, (mlp_output, self.mlp_sconv)
230
+ return _sconv_add_norm(
231
+ mlp_output, hidden_states, self.mlp_sconv, None, positions
232
+ )[1]
233
234
235
class InklingReplicatedEmbedding(nn.Module):
@@ -330,6 +338,7 @@ class InklingModel(nn.Module):
338
positions,
339
hidden_states,
340
pending=pending,
341
+ defer_mlp_add=True,
342
attn_in=attn_in0,
343
log_scaling=log_scaling,
344
)
vllm/models/inkling/nvidia/mtp.py
new
+408
@@ -0,0 +1,408 @@
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
+Implements the first MTP depth from the reference ``mtp_model.py`` shipped with
6
+the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H``
7
+input projection, and a full Inkling transformer block with a dense bf16 MLP.
8
+
9
+The draft shares the target's token embedding table and LM head
10
+(``load_eagle_model`` wires those references) and applies the backbone
11
+``embed_norm`` on top: the depth layers were trained on the same normed
12
+embeddings the backbone consumes (their own ``embed_norm`` weights are
13
+near-identity trims, unlike the backbone's whitening ``embed_norm``).
14
+"""
15
+
16
+from __future__ import annotations
17
+
18
+from collections.abc import Iterable
19
+
20
+import regex as re
21
+import torch
22
+from torch import nn
23
+
24
+from vllm.config import VllmConfig
25
+from vllm.model_executor.layers.linear import ReplicatedLinear
26
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
28
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29
+from vllm.model_executor.models.utils import maybe_prefix
30
+from vllm.sequence import IntermediateTensors
31
+
32
+from ..configs import InklingModelConfig
33
+from .layernorm import InklingRMSNorm
34
+from .model import InklingDecoderLayer, InklingReplicatedEmbedding
35
+from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm
36
+
37
+# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr.
38
+# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a
39
+# local (pname, wname, shard) list since the MTP loader remaps by hand.
40
+_ATTENTION_PARAMS_MAPPING = [
41
+ ("qkvr", "wq_du", 0),
42
+ ("qkvr", "wk_dv", 1),
43
+ ("qkvr", "wv_dv", 2),
44
+ ("qkvr", "wr_du", 3),
45
+]
46
+
47
+
48
+def _mtp_depth_from_name(name: str) -> int | None:
49
+ m = re.search(r"\.mtp\.layers\.(\d+)\.", name)
50
+ return int(m.group(1)) if m else None
51
+
52
+
53
+class InklingMTPDepthLayer(nn.Module):
54
+ """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block."""
55
+
56
+ def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None:
57
+ super().__init__()
58
+ self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
59
+ self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
60
+ self.input_proj = ReplicatedLinear(
61
+ config.hidden_size * 2,
62
+ config.hidden_size,
63
+ bias=False,
64
+ return_bias=False,
65
+ prefix=f"{prefix}.input_proj",
66
+ )
67
+ # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs
68
+ # full attention (the swa_* head config and sliding_window window) to
69
+ # match this depth's checkpoint transformer_block weights.
70
+ self.transformer_block = InklingDecoderLayer(
71
+ config,
72
+ layer_id=0,
73
+ is_local=is_local,
74
+ quant_config=None,
75
+ prefix=f"{prefix}.transformer_block",
76
+ nvfp4_config=None,
77
+ force_dense_mlp=True,
78
+ )
79
+
80
+ def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
81
+ # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)]
82
+ # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch.
83
+ hidden = self.input_proj(combined)
84
+ # The short conv self-fetches its paged SWA-cache metadata from the
85
+ # forward context (via its conv_owner prefix); no conv_meta to thread.
86
+ return self.transformer_block(positions, hidden)
87
+
88
+
89
+class InklingMultiTokenPredictor(nn.Module):
90
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
91
+ super().__init__()
92
+ assert vllm_config.speculative_config is not None
93
+ config: InklingModelConfig = (
94
+ vllm_config.speculative_config.draft_model_config.hf_config
95
+ )
96
+ self.config = config
97
+ if vllm_config.speculative_config.num_speculative_tokens != 1:
98
+ raise ValueError(
99
+ "Inkling MTP currently supports exactly one speculative token"
100
+ )
101
+ self.chain_hidden_post_norm = config.chain_hidden_post_norm
102
+ local_ids = set(config.local_layer_ids)
103
+ self.layers = nn.ModuleDict(
104
+ {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)}
105
+ )
106
+ self.chain_norm = (
107
+ InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
108
+ if self.chain_hidden_post_norm
109
+ else None
110
+ )
111
+ # The target's raw token embedding (pre embed_norm), attached by
112
+ # load_eagle_model. Never materialized here: building our own
113
+ # replicated copy would transiently double the 2.3 GiB table.
114
+ self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment]
115
+ # The depth layers consume the *backbone-normed* embedding
116
+ # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights
117
+ # are near-identity (trained on already-normalized inputs), and
118
+ # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70.
119
+ # Weight loaded from the target's embed_norm.weight; gated like the
120
+ # target's InklingModel.embed_norm.
121
+ self.backbone_embed_norm = (
122
+ InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
123
+ if config.use_embed_norm
124
+ else None
125
+ )
126
+
127
+ def embed_input_ids(
128
+ self,
129
+ input_ids: torch.Tensor,
130
+ multimodal_embeddings: object | None = None,
131
+ *,
132
+ is_multimodal: torch.Tensor | None = None,
133
+ ) -> torch.Tensor:
134
+ """Draft-prefill embedding: fused gather + backbone embed_norm, then
135
+ the target's tower embeddings scattered in unnormed (the backbone
136
+ convention — MM embeds are merged after embed_norm)."""
137
+ norm = self.backbone_embed_norm
138
+ embeds = embed_rmsnorm(
139
+ input_ids,
140
+ self.embed_tokens.weight,
141
+ norm.weight if norm is not None else None,
142
+ norm.variance_epsilon if norm is not None else 0.0,
143
+ )
144
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type]
145
+ return embeds
146
+ from vllm.model_executor.models.utils import _merge_multimodal_embeddings
147
+
148
+ assert is_multimodal is not None
149
+ return _merge_multimodal_embeddings(
150
+ inputs_embeds=embeds,
151
+ multimodal_embeddings=multimodal_embeddings,
152
+ is_multimodal=is_multimodal,
153
+ )
154
+
155
+ def fused_input_cat(
156
+ self,
157
+ layer: InklingMTPDepthLayer,
158
+ previous_hidden: torch.Tensor,
159
+ input_ids: torch.Tensor,
160
+ inputs_embeds: torch.Tensor | None,
161
+ ) -> torch.Tensor:
162
+ """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one
163
+ launch: embedding row gather + the backbone embed_norm + the depth
164
+ embed_norm chain on one side, hidden_norm on the other, written
165
+ straight into the cat buffer."""
166
+ hidden_w = layer.hidden_norm.weight
167
+ embed_w = layer.embed_norm.weight
168
+ eps = layer.hidden_norm.variance_epsilon
169
+ if inputs_embeds is not None:
170
+ # Draft prefill with target-merged MM embeddings (already
171
+ # backbone-normed via embed_input_ids); only the depth embed_norm
172
+ # remains.
173
+ return embed_dual_rmsnorm_cat(
174
+ previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds
175
+ )
176
+ return embed_dual_rmsnorm_cat(
177
+ previous_hidden,
178
+ hidden_w,
179
+ embed_w,
180
+ eps,
181
+ input_ids=input_ids,
182
+ embed_table=self.embed_tokens.weight,
183
+ pre_norm_weight=(
184
+ self.backbone_embed_norm.weight
185
+ if self.backbone_embed_norm is not None
186
+ else None
187
+ ),
188
+ )
189
+
190
+ def forward(
191
+ self,
192
+ input_ids: torch.Tensor,
193
+ positions: torch.Tensor,
194
+ previous_hidden_states: torch.Tensor,
195
+ inputs_embeds: torch.Tensor | None = None,
196
+ spec_step_idx: int = 0,
197
+ ) -> torch.Tensor:
198
+ # The draft's short conv is a paged SWA-cache layer (its conv_owner is
199
+ # auto-enumerated as a draft attention layer); its per-token metadata is
200
+ # built by the speculator's build_attn_metadata and read from the
201
+ # forward context, so nothing extra is threaded here.
202
+ if spec_step_idx != 0:
203
+ raise ValueError("Inkling MTP only supports spec_step_idx=0")
204
+ layer = self.layers["0"]
205
+ combined = self.fused_input_cat(
206
+ layer, previous_hidden_states, input_ids, inputs_embeds
207
+ )
208
+ hidden = layer(combined, positions)
209
+ if self.chain_norm is not None:
210
+ hidden = self.chain_norm(hidden)
211
+ return hidden
212
+
213
+
214
+class InklingMTP(nn.Module):
215
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
216
+ super().__init__()
217
+ assert vllm_config.speculative_config is not None
218
+ config: InklingModelConfig = (
219
+ vllm_config.speculative_config.draft_model_config.hf_config
220
+ )
221
+ self.config = config
222
+ self.model = InklingMultiTokenPredictor(
223
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
224
+ )
225
+ # The target's (vocab-sharded) LM head, attached by load_eagle_model;
226
+ # never materialized here (same reasoning as model.embed_tokens).
227
+ self.lm_head: ParallelLMHead = None # type: ignore[assignment]
228
+ self.logits_processor = LogitsProcessor(
229
+ config.padded_vocab_size,
230
+ org_vocab_size=config.vocab_size,
231
+ soft_cap=config.final_logit_softcapping,
232
+ )
233
+ self._logits_zero: torch.Tensor | None = None
234
+
235
+ def embed_input_ids(
236
+ self,
237
+ input_ids: torch.Tensor,
238
+ multimodal_embeddings: object | None = None,
239
+ *,
240
+ is_multimodal: torch.Tensor | None = None,
241
+ ) -> torch.Tensor:
242
+ return self.model.embed_input_ids(
243
+ input_ids, multimodal_embeddings, is_multimodal=is_multimodal
244
+ )
245
+
246
+ def forward(
247
+ self,
248
+ input_ids: torch.Tensor,
249
+ positions: torch.Tensor,
250
+ hidden_states: torch.Tensor,
251
+ intermediate_tensors: IntermediateTensors | None = None,
252
+ inputs_embeds: torch.Tensor | None = None,
253
+ spec_step_idx: int = 0,
254
+ ) -> torch.Tensor:
255
+ return self.model(
256
+ input_ids,
257
+ positions,
258
+ hidden_states,
259
+ inputs_embeds,
260
+ spec_step_idx,
261
+ )
262
+
263
+ def compute_logits(
264
+ self,
265
+ hidden_states: torch.Tensor,
266
+ spec_step_idx: int = 0,
267
+ ) -> torch.Tensor | None:
268
+ # The MTP shares the base model's LM head, which is trained on
269
+ # ``hidden / mup``-scaled inputs, so apply the same mup scaling here
270
+ # for a matching logit scale — folded into the lm_head GEMM alpha
271
+ # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant
272
+ # for greedy draft sampling, but it matters for the gumbel sampling
273
+ # distribution at temperature > 0.)
274
+ mup = self.config.logits_mup_width_multiplier
275
+ if not mup:
276
+ return self.logits_processor(self.lm_head, hidden_states)
277
+ assert self.logits_processor.soft_cap is None
278
+ assert self.logits_processor.scale == 1.0
279
+ w = self.lm_head.weight
280
+ if self._logits_zero is None:
281
+ self._logits_zero = w.new_zeros(1)
282
+ logits = torch.addmm(
283
+ self._logits_zero,
284
+ hidden_states,
285
+ w.t(),
286
+ beta=0.0,
287
+ alpha=1.0 / mup,
288
+ )
289
+ logits = self.logits_processor._gather_logits(logits)
290
+ if logits is not None:
291
+ logits = logits[..., : self.logits_processor.org_vocab_size]
292
+ return logits
293
+
294
+ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
295
+ """Greedy draft tokens via rank-local argmax + tiny (value, index)
296
+ reduction — no full-vocab logits all-gather. The muP divisor is a
297
+ positive scalar, so the argmax is invariant and the scaling is
298
+ skipped entirely."""
299
+ return self.logits_processor.get_top_tokens(self.lm_head, hidden_states)
300
+
301
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
302
+ return _load_inkling_mtp_weights(self, weights)
303
+
304
+
305
+def _load_inkling_mtp_weights(
306
+ module: InklingMTP,
307
+ weights: Iterable[tuple[str, torch.Tensor]],
308
+) -> set[str]:
309
+ """Load ``model.mtp.*`` weights into the MTP module.
310
+
311
+ Checkpoint keys look like ``model.mtp.chain_norm.weight`` and
312
+ ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone
313
+ layer's fused-projection layout, so we apply the same qkvr / gate_up / down
314
+ remapping as ``_load_inkling_weights``. Token embedding and LM head are shared
315
+ (provided by ``load_eagle_model``) and are not present in mtp.safetensors.
316
+ """
317
+ # Per-depth attention is full or sliding-window (config.local_layer_ids);
318
+ # each depth's qkvr MergedColumnParallelLinear is built with the matching
319
+ # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding.
320
+ # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never
321
+ # exceeds a layer's kv-head count and no GQA K/V replication is needed here.
322
+ params = dict(module.named_parameters())
323
+ loaded: set[str] = set()
324
+
325
+ def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool:
326
+ param = params.get(name)
327
+ if param is None:
328
+ return False
329
+ loader = getattr(param, "weight_loader", default_weight_loader)
330
+ if shard_id is None:
331
+ if loader is default_weight_loader or param.shape == weight.shape:
332
+ default_weight_loader(param, weight)
333
+ else:
334
+ loader(param, weight)
335
+ else:
336
+ loader(param, weight, shard_id) # type: ignore[call-arg]
337
+ loaded.add(name)
338
+ return True
339
+
340
+ for name, weight in weights:
341
+ depth = _mtp_depth_from_name(name)
342
+ # Token embedding and LM head are never materialized on the draft
343
+ # (no params to load into); load_eagle_model attaches the target's.
344
+ if name in ("model.llm.embed.weight", "model.llm.unembed.weight"):
345
+ continue
346
+ # The backbone embed_norm, applied to the shared embedding before the
347
+ # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The
348
+ # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded
349
+ # below. Only the shared backbone key routes here.
350
+ if name == "model.llm.embed_norm.weight":
351
+ _load("model.backbone_embed_norm.weight", weight)
352
+ continue
353
+ # Only consume the MTP weights; everything else belongs to the target.
354
+ if ".mtp." not in name:
355
+ continue
356
+ # Only the first checkpoint depth is used for MTP=1.
357
+ if depth is not None and depth != 0:
358
+ continue
359
+ # model.mtp.chain_norm.weight -> model.chain_norm.weight
360
+ # model.mtp.layers.{i}.X -> model.layers.{i}.X
361
+ original_name = name
362
+ name = name.replace(".mtp.layers.", ".layers.").replace(
363
+ ".mtp.chain_norm.", ".chain_norm."
364
+ )
365
+
366
+ if ".chain_norm." in name and module.model.chain_norm is None:
367
+ raise ValueError(
368
+ "Inkling checkpoint contains chain_norm weights but "
369
+ "chain_hidden_post_norm is disabled."
370
+ )
371
+
372
+ # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr).
373
+ matched = False
374
+ for pname, wname, shard in _ATTENTION_PARAMS_MAPPING:
375
+ if f".attn.{wname}." in name:
376
+ mapped_name = name.replace(f".{wname}.", f".{pname}.")
377
+ if not _load(mapped_name, weight, shard):
378
+ raise ValueError(f"Unexpected Inkling MTP weight: {original_name}")
379
+ matched = True
380
+ break
381
+ if matched:
382
+ continue
383
+
384
+ # Dense MLP fused gate/up + down.
385
+ if ".mlp.w13_dn.weight" in name:
386
+ loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight)
387
+ elif ".mlp.w2_md.weight" in name:
388
+ loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight)
389
+ else:
390
+ if name.endswith(".bias") and name not in params:
391
+ continue
392
+ loaded_weight = _load(name, weight)
393
+ if not loaded_weight:
394
+ raise ValueError(f"Unexpected Inkling MTP weight: {original_name}")
395
+ required = {
396
+ name
397
+ for name in params
398
+ if name.startswith("model.layers.") or name.startswith("model.chain_norm.")
399
+ }
400
+ if missing := sorted(required - loaded):
401
+ raise ValueError(
402
+ "Inkling MTP checkpoint is missing required parameters: "
403
+ + ", ".join(missing)
404
+ )
405
+ return loaded
406
+
407
+
408
+EntryClass = [InklingMTP]
vllm/models/inkling/nvidia/ops/norm.py
+99
@@ -268,6 +268,105 @@ def embed_rmsnorm(
268
return out
269
270
271
+@triton.jit
272
+def _embed_dual_rmsnorm_cat_kernel(
273
+ hidden_ptr, # [T, N]
274
+ emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER
275
+ ids_ptr, # [T] token ids (GATHER only)
276
+ w_hidden_ptr, # [N]
277
+ w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only)
278
+ w_embed_ptr, # [N]
279
+ out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))]
280
+ eps,
281
+ hidden_stride_0,
282
+ emb_stride_0,
283
+ n_cols,
284
+ block_size_n: tl.constexpr,
285
+ GATHER: tl.constexpr,
286
+ HAS_PRE_NORM: tl.constexpr,
287
+):
288
+ pid_m = tl.program_id(0).to(tl.int64)
289
+ which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N)
290
+ offs_n = tl.arange(0, block_size_n)
291
+ mask_n = offs_n < n_cols
292
+ if which == 0:
293
+ x = tl.load(
294
+ hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0
295
+ ).to(tl.float32)
296
+ w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
297
+ else:
298
+ row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m
299
+ x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to(
300
+ tl.float32
301
+ )
302
+ if HAS_PRE_NORM:
303
+ w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
304
+ rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps)
305
+ # Round-trip through the output dtype so the chained norm is
306
+ # bit-exact vs the unfused pair (which stores bf16 in between).
307
+ x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32)
308
+ w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
309
+ rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps)
310
+ tl.store(
311
+ out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n,
312
+ (x * rstd * w).to(out_ptr.dtype.element_ty),
313
+ mask=mask_n,
314
+ )
315
+
316
+
317
+def embed_dual_rmsnorm_cat(
318
+ hidden: torch.Tensor,
319
+ hidden_weight: torch.Tensor,
320
+ embed_weight: torch.Tensor,
321
+ eps: float,
322
+ *,
323
+ embeds: torch.Tensor | None = None,
324
+ input_ids: torch.Tensor | None = None,
325
+ embed_table: torch.Tensor | None = None,
326
+ pre_norm_weight: torch.Tensor | None = None,
327
+) -> torch.Tensor:
328
+ """The MTP depth-layer input in one launch:
329
+ ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``.
330
+
331
+ The embed side is either a fused row gather ``embed_table[input_ids]``
332
+ (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged
333
+ multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the
334
+ backbone embed_norm in front of the depth embed_norm (bit-exact vs the
335
+ unfused sequence). The concat copies collapse into direct writes."""
336
+ T, n = hidden.shape
337
+ if embeds is not None:
338
+ assert embeds.shape == hidden.shape
339
+ src, ids, src_stride = embeds, embeds, embeds.stride(0)
340
+ gather = False
341
+ else:
342
+ assert input_ids is not None and embed_table is not None
343
+ assert input_ids.shape == (T,) and embed_table.shape[1] == n
344
+ src, ids, src_stride = embed_table, input_ids, embed_table.stride(0)
345
+ gather = True
346
+ out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device)
347
+ if T == 0:
348
+ return out
349
+ block_size_n = triton.next_power_of_2(n)
350
+ _embed_dual_rmsnorm_cat_kernel[(T, 2)](
351
+ hidden,
352
+ src,
353
+ ids,
354
+ hidden_weight,
355
+ pre_norm_weight if pre_norm_weight is not None else embed_weight,
356
+ embed_weight,
357
+ out,
358
+ eps,
359
+ hidden.stride(0),
360
+ src_stride,
361
+ n,
362
+ block_size_n,
363
+ GATHER=gather,
364
+ HAS_PRE_NORM=pre_norm_weight is not None,
365
+ num_warps=_get_num_warps_from_block_size(block_size_n),
366
+ )
367
+ return out
368
+
369
+
370
def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
371
assert x.ndim == 2, f"{x.shape=}"
372
assert weight.ndim == 1, f"{weight.shape=}"