[ROCm] Restore Inkling MTP backend parity (#50806)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 4, 2026 at 21:02 UTC
33c50587d2679ba9bacc2a51ae19901f7eb3a129
1 file changed
+42
-13
vllm/models/inkling/amd/mtp.py
+42
-13
@@ -2,9 +2,13 @@
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.
5
+Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP
6
+depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj``
7
+(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the
8
+same short convolutions as the backbone; its attention is full or sliding-window
9
+per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared
10
+``chain_norm`` is applied after every depth; its output is both the logits input
11
+and the previous hidden state fed to the next depth.
12
13
The draft shares the target's token embedding table and LM head
14
(``load_eagle_model`` wires those references) and applies the backbone
@@ -54,6 +58,16 @@ def _mtp_depth_from_name(name: str) -> int | None:
58
return int(m.group(1)) if m else None
59
60
61
+def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int:
62
+ num_layers = min(n_predict, num_spec) if num_spec else n_predict
63
+ if num_layers <= 0:
64
+ raise ValueError(
65
+ "Inkling MTP requires num_nextn_predict_layers and "
66
+ "num_speculative_tokens to select at least one depth layer."
67
+ )
68
+ return num_layers
69
+
70
+
71
class InklingMTPDepthLayer(nn.Module):
72
"""One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block."""
73
@@ -97,14 +111,30 @@ class InklingMultiTokenPredictor(nn.Module):
111
vllm_config.speculative_config.draft_model_config.hf_config
112
)
113
self.config = config
100
- if vllm_config.speculative_config.num_speculative_tokens != 1:
101
- raise ValueError(
102
- "Inkling MTP currently supports exactly one speculative token"
103
- )
114
+ # The checkpoint ships num_nextn_predict_layers depth blocks, but only
115
+ # the first ``num_speculative_tokens`` are exercised (step i uses depth
116
+ # i). Build only those to save memory — each depth is a full Inkling block
117
+ # with its own (large) full-history sconv caches and KV cache.
118
+ n_predict = config.num_nextn_predict_layers
119
+ num_spec = vllm_config.speculative_config.num_speculative_tokens
120
+ self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec)
121
self.chain_hidden_post_norm = config.chain_hidden_post_norm
122
+
123
+ # Depth blocks whose attention is sliding-window (swa_* head config)
124
+ # rather than full; keyed by MTP depth via the checkpoint's
125
+ # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors
126
+ # InklingModel's local_ids split, but over MTP depths, not backbone
127
+ # layers.
128
local_ids = set(config.local_layer_ids)
129
+
130
+ # Keyed by depth index (str) to mirror the checkpoint layout.
131
self.layers = nn.ModuleDict(
107
- {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)}
132
+ {
133
+ str(idx): InklingMTPDepthLayer(
134
+ config, f"{prefix}.layers.{idx}", idx in local_ids
135
+ )
136
+ for idx in range(self.num_mtp_layers)
137
+ }
138
)
139
self.chain_norm = (
140
InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -202,9 +232,8 @@ class InklingMultiTokenPredictor(nn.Module):
232
# auto-enumerated as a draft attention layer); its per-token metadata is
233
# built by the speculator's build_attn_metadata and read from the
234
# forward context, so nothing extra is threaded here.
205
- if spec_step_idx != 0:
206
- raise ValueError("Inkling MTP only supports spec_step_idx=0")
207
- layer = self.layers["0"]
235
+ depth = spec_step_idx % self.num_mtp_layers
236
+ layer = self.layers[str(depth)]
237
combined = self.fused_input_cat(
238
layer, previous_hidden_states, input_ids, inputs_embeds
239
)
@@ -356,8 +385,8 @@ def _load_inkling_mtp_weights(
385
# Only consume the MTP weights; everything else belongs to the target.
386
if ".mtp." not in name:
387
continue
359
- # Only the first checkpoint depth is used for MTP=1.
360
- if depth is not None and depth != 0:
388
+ # Skip depth blocks beyond the ones we built (num_speculative_tokens).
389
+ if depth is not None and depth >= module.model.num_mtp_layers:
390
continue
391
# model.mtp.chain_norm.weight -> model.chain_norm.weight
392
# model.mtp.layers.{i}.X -> model.layers.{i}.X