[Model Runner V2][Spec Decode] Add multi-layer MTP speculator (#48892)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

vllm/config/speculative.py

@@ -1410,6 +1410,14 @@ class SpeculativeConfig:
     def use_ngram_gpu(self) -> bool:
         return self.method == "ngram_gpu"
 
+    def use_multi_module_mtp(self) -> bool:
+        if self.method != "mtp" or self.draft_model_config is None:
+            return False
+        num_mtp_layers = getattr(
+            self.draft_model_config.hf_config, "num_nextn_predict_layers", 1
+        )
+        return min(num_mtp_layers, self.num_speculative_tokens) > 1
+
     def __repr__(self) -> str:
         method = self.method
         model = (

vllm/models/inkling/nvidia/mtp.py

@@ -2,9 +2,13 @@
 # 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.
+Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP
+depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj``
+(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the
+same short convolutions as the backbone; its attention is full or sliding-window
+per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared
+``chain_norm`` is applied after every depth; its output is both the logits input
+and the previous hidden state fed to the next depth.
 
 The draft shares the target's token embedding table and LM head
 (``load_eagle_model`` wires those references) and applies the backbone
@@ -53,6 +57,16 @@ def _mtp_depth_from_name(name: str) -> int | None:
     return int(m.group(1)) if m else None
 
 
+def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int:
+    num_layers = min(n_predict, num_spec) if num_spec else n_predict
+    if num_layers <= 0:
+        raise ValueError(
+            "Inkling MTP requires num_nextn_predict_layers and "
+            "num_speculative_tokens to select at least one depth layer."
+        )
+    return num_layers
+
+
 class InklingMTPDepthLayer(nn.Module):
     """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block."""
 
@@ -96,14 +110,30 @@ class InklingMultiTokenPredictor(nn.Module):
             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"
-            )
+        # The checkpoint ships num_nextn_predict_layers depth blocks, but only
+        # the first ``num_speculative_tokens`` are exercised (step i uses depth
+        # i). Build only those to save memory — each depth is a full Inkling block
+        # with its own (large) full-history sconv caches and KV cache.
+        n_predict = config.num_nextn_predict_layers
+        num_spec = vllm_config.speculative_config.num_speculative_tokens
+        self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec)
         self.chain_hidden_post_norm = config.chain_hidden_post_norm
+
+        # Depth blocks whose attention is sliding-window (swa_* head config)
+        # rather than full; keyed by MTP depth via the checkpoint's
+        # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors
+        # InklingModel's local_ids split, but over MTP depths, not backbone
+        # layers.
         local_ids = set(config.local_layer_ids)
+
+        # Keyed by depth index (str) to mirror the checkpoint layout.
         self.layers = nn.ModuleDict(
-            {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)}
+            {
+                str(idx): InklingMTPDepthLayer(
+                    config, f"{prefix}.layers.{idx}", idx in local_ids
+                )
+                for idx in range(self.num_mtp_layers)
+            }
         )
         self.chain_norm = (
             InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -201,9 +231,8 @@ class InklingMultiTokenPredictor(nn.Module):
         # 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"]
+        depth = spec_step_idx % self.num_mtp_layers
+        layer = self.layers[str(depth)]
         combined = self.fused_input_cat(
             layer, previous_hidden_states, input_ids, inputs_embeds
         )
@@ -355,8 +384,8 @@ def _load_inkling_mtp_weights(
         # 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:
+        # Skip depth blocks beyond the ones we built (num_speculative_tokens).
+        if depth is not None and depth >= module.model.num_mtp_layers:
             continue
         # model.mtp.chain_norm.weight -> model.chain_norm.weight
         # model.mtp.layers.{i}.X      -> model.layers.{i}.X

vllm/v1/worker/gpu/input_batch.py

@@ -191,6 +191,8 @@ class InputBatch:
 def _prepare_prefill_inputs_kernel(
     input_ids_ptr,
     next_prefill_tokens_ptr,
+    next_prefill_tokens_stride,
+    num_lookahead,
     idx_mapping_ptr,
     query_start_loc_ptr,
     all_token_ids_ptr,
@@ -198,6 +200,7 @@ def _prepare_prefill_inputs_kernel(
     prefill_lens_ptr,
     num_computed_tokens_ptr,
     BLOCK_SIZE: tl.constexpr,
+    LOOKAHEAD_BLOCK: tl.constexpr,
 ):
     batch_idx = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
@@ -218,10 +221,20 @@ def _prepare_prefill_inputs_kernel(
         tokens = tl.load(request_ptr + num_computed + block, mask=mask)
         tl.store(input_ids_ptr + query_start + block, tokens, mask=mask)
 
-    next_pos = num_computed + query_len
-    if next_pos < prefill_len:
-        next_token = tl.load(request_ptr + next_pos)
-        tl.store(next_prefill_tokens_ptr + req_state_idx, next_token)
+    # Store the next num_lookahead prefill tokens.
+    lookahead = tl.arange(0, LOOKAHEAD_BLOCK)
+    pos = num_computed + query_len + lookahead
+    in_lookahead = lookahead < num_lookahead
+    tokens = tl.load(
+        request_ptr + pos, mask=in_lookahead & (pos < prefill_len), other=0
+    )
+    tl.store(
+        next_prefill_tokens_ptr
+        + lookahead * next_prefill_tokens_stride
+        + req_state_idx,
+        tokens,
+        mask=in_lookahead,
+    )
 
 
 def prepare_prefill_inputs(
@@ -234,9 +247,12 @@ def prepare_prefill_inputs(
     num_computed_tokens: torch.Tensor,
 ) -> None:
     num_reqs = idx_mapping.shape[0]
+    num_lookahead = next_prefill_tokens.shape[0]
     _prepare_prefill_inputs_kernel[(num_reqs,)](
         input_ids,
         next_prefill_tokens,
+        next_prefill_tokens.stride(0),
+        num_lookahead,
         idx_mapping,
         query_start_loc,
         all_token_ids,
@@ -244,6 +260,7 @@ def prepare_prefill_inputs(
         prefill_len,
         num_computed_tokens,
         BLOCK_SIZE=1024,
+        LOOKAHEAD_BLOCK=triton.next_power_of_2(num_lookahead),
     )
 
 

vllm/v1/worker/gpu/model_runner.py

@@ -223,6 +223,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.is_pooling_model = self.model_config.runner_type == "pooling"
         self.pooling_runner: PoolingRunner | None = None
 
+        # Multi-module MTP feeds its modules the next num_speculative_steps prefill
+        # tokens during chunked prefill. Other speculators only read the immediate
+        # next one.
+        num_prefill_lookahead = (
+            self.num_speculative_steps
+            if self.speculative_config is not None
+            and self.speculative_config.use_multi_module_mtp()
+            else 1
+        )
         # General request states.
         self.req_states = RequestState(
             max_num_reqs=self.max_num_reqs,
@@ -231,6 +240,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_speculative_steps=self.num_speculative_steps,
             vocab_size=self.vocab_size,
             device=self.device,
+            num_prefill_lookahead=num_prefill_lookahead,
         )
         self.input_buffers = InputBuffers(
             max_num_reqs=self.max_num_reqs,
@@ -628,7 +638,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     torch.zeros(
                         input_batch.num_tokens,
                         dtype=torch.bool,
-                        device=self.device,
+                        device="cpu",
                     ),
                 )
 
@@ -1521,6 +1531,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # NOTE: This is done here because postprocess updates
             # num_computed_prefill_tokens.
             # The EAGLE/MTP drafter reads one position ahead of the target.
+            # TODO(TheEpicDolphin): Gather MM embeddings for all speculative
+            # steps during multi-module MTP.
             mm_inputs = self.model_state.gather_mm_embeddings(
                 input_batch, draft_lookahead=1
             )

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 from typing import Any
 
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -282,6 +283,7 @@ class DFlashSpeculator(DraftModelSpeculator):
         step: int,
         num_query_per_req: int | None = None,
         causal: bool | Mapping[int, bool] = False,
+        query_start_loc_np: np.ndarray | None = None,
     ) -> dict[str, Any] | None:
         if not self.draft_attn_layer_names:
             return None
@@ -294,6 +296,7 @@ class DFlashSpeculator(DraftModelSpeculator):
             step=step,
             num_query_per_req=self.num_query_per_req,
             causal=causal,
+            query_start_loc_np=query_start_loc_np,
         )
 
     @torch.inference_mode()

vllm/v1/worker/gpu/spec_decode/multi_module_mtp/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from typing import Any
 
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -51,7 +52,7 @@ class BaseSpeculator(ABC):
         num_rejected: torch.Tensor,
         # [max_num_reqs]
         last_sampled: torch.Tensor,
-        # [max_num_reqs]
+        # [num_prefill_lookahead, max_num_reqs]
         next_prefill_tokens: torch.Tensor,
         # [max_num_reqs]
         temperature: torch.Tensor,
@@ -214,14 +215,28 @@ class DraftModelSpeculator(BaseSpeculator):
         step: int,
         num_query_per_req: int = 1,
         causal: bool | Mapping[int, bool] = True,
+        query_start_loc_np: np.ndarray | None = None,
     ) -> dict[str, Any] | None:
-        # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
-        # Clamp keeps the series non-decreasing past num_reqs, which some
-        # attention backends require.
-        query_start_loc_cpu = (
-            torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
-            * num_query_per_req
-        )
+        if query_start_loc_np is not None:
+            # Non-uniform query layout (e.g. multi-module MTP's mixed
+            # prefill/decode queries); num_query_per_req is ignored.
+            query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
+            query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
+                query_start_loc_np[: num_reqs + 1]
+            )
+            query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
+            max_query_len = int(
+                (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()
+            )
+        else:
+            # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
+            # Clamp keeps the series non-decreasing past num_reqs, which some
+            # attention backends require.
+            query_start_loc_cpu = (
+                torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
+                * num_query_per_req
+            )
+            max_query_len = num_query_per_req
         block_tables = [
             x[:num_reqs_padded] for x in self.block_tables.input_block_tables
         ]
@@ -243,7 +258,7 @@ class DraftModelSpeculator(BaseSpeculator):
                 : num_reqs_padded + 1
             ],
             query_start_loc_cpu=query_start_loc_cpu,
-            max_query_len=num_query_per_req,
+            max_query_len=max_query_len,
             seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
             max_seq_len=self.draft_max_seq_len,
             block_tables=block_tables,
@@ -325,7 +340,6 @@ class DraftModelSpeculator(BaseSpeculator):
         self.temperature.copy_(temperature)
         self.seeds.copy_(seeds)
         self.idx_mapping[:num_reqs].copy_(idx_mapping)
-        if self.draft_logits is not None:
-            # idx_mapping for CG padded requests points to -1, which is ignored
-            # during sampling to prevent writing stale values to draft logits.
-            self.idx_mapping[num_reqs:].fill_(-1)
+        # idx_mapping for CG padded requests points to -1, which is ignored
+        # during sampling to prevent writing stale values to draft logits.
+        self.idx_mapping[num_reqs:].fill_(-1)

vllm/v1/worker/gpu/states.py

@@ -15,6 +15,7 @@ class RequestState:
         num_speculative_steps: int,
         vocab_size: int,
         device: torch.device,
+        num_prefill_lookahead: int = 1,
     ):
         self.max_num_reqs = max_num_reqs
         self.max_model_len = max_model_len
@@ -77,7 +78,10 @@ class RequestState:
         )
 
         self.next_prefill_tokens = torch.zeros(
-            self.max_num_reqs, dtype=torch.int32, device=device
+            num_prefill_lookahead,
+            self.max_num_reqs,
+            dtype=torch.int32,
+            device=device,
         )
 
     @property