Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin

nemo/collections/speechlm2/vllm/salm/__init__.py

@@ -26,6 +26,84 @@
 _PKG = "nemo.collections.speechlm2.vllm.salm"
 
 
+def _patch_vllm_for_nemo_speechlm_mtp() -> None:
+    """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
+
+    Three patches are applied:
+
+    1. ``MTPModelTypes`` — the Literal type that guards the MTP detection
+       branch in ``SpeculativeConfig.__post_init__`` is extended to include
+       ``"nemo_speechlm_mtp"``.
+
+    2. ``SpeculativeConfig.hf_config_override`` — the static method that
+       rewrites the draft-model HF config is wrapped to detect
+       ``nemo_speechlm`` checkpoints that carry MTP heads (``mtp.enabled``
+       and ``mtp.num_nextn_predict_layers > 0``) and redirect them to the
+       ``NeMoSpeechLMMTPModel`` architecture with the right ``n_predict``.
+
+    3. ``ModelRegistry`` — ``NeMoSpeechLMMTPModel`` is registered so that
+       vLLM can resolve and instantiate it as the draft model.
+    """
+    from typing import Literal, get_args
+
+    import vllm.config.speculative as _spec_mod
+    from vllm.config.speculative import SpeculativeConfig
+
+    # 1 ── Extend MTPModelTypes -------------------------------------------
+    old_args = get_args(_spec_mod.MTPModelTypes)
+    if "nemo_speechlm_mtp" not in old_args:
+        _spec_mod.MTPModelTypes = Literal[old_args + ("nemo_speechlm_mtp",)]
+
+    # 2 ── Wrap hf_config_override ----------------------------------------
+    current_override = SpeculativeConfig.hf_config_override
+    if not getattr(current_override, "_nemo_speechlm_mtp_override", False):
+        original_override = current_override
+
+        def _patched_override(hf_config):
+            if hf_config.model_type == "nemo_speechlm":
+                mtp_cfg = getattr(hf_config, "mtp", {}) or {}
+                n_predict = mtp_cfg.get("num_nextn_predict_layers", 0) if isinstance(mtp_cfg, dict) else 0
+                if n_predict > 0:
+                    use_repeated_layer = bool(mtp_cfg.get("use_repeated_layer", False))
+                    if n_predict > 1 and not use_repeated_layer:
+                        raise ValueError(
+                            f"NeMo SpeechLM MTP with {n_predict} distinct head layers is not "
+                            f"supported: vLLM's NemotronHMultiTokenPredictor builds a single "
+                            f"physical MTP layer and reuses it every speculative step. Only "
+                            f"checkpoints trained with mtp.use_repeated_layer=true match that "
+                            f"execution model."
+                        )
+                    hf_config.model_type = "nemo_speechlm_mtp"
+                    hf_config.update(
+                        {
+                            # Draft-able tokens per target step: SpeculativeConfig caps/validates
+                            # num_speculative_tokens against this.
+                            "n_predict": n_predict,
+                            # Physical MTP layers to instantiate. Repeated-layer checkpoints ship
+                            # one shared layer (mtp.layers.0.*) that is reapplied every step, which
+                            # is exactly how the vLLM proposer drives an MTP draft. This also
+                            # shadows the backbone text_config's num_nextn_predict_layers (e.g. 4),
+                            # which would otherwise trip the single-layer assert in
+                            # NemotronHMultiTokenPredictor.
+                            "num_nextn_predict_layers": 1,
+                            "architectures": ["NeMoSpeechLMMTPModel"],
+                        }
+                    )
+                    return hf_config
+            return original_override(hf_config)
+
+        _patched_override._nemo_speechlm_mtp_override = True
+        SpeculativeConfig.hf_config_override = staticmethod(_patched_override)
+
+    # 3 ── Register NeMoSpeechLMMTPModel ----------------------------------
+    from vllm.model_executor.models.registry import ModelRegistry
+
+    ModelRegistry.register_model(
+        "NeMoSpeechLMMTPModel",
+        f"{_PKG}.mtp:NeMoSpeechLMMTP",
+    )
+
+
 def register():
     """Register the NeMo Speech LM model and config with vLLM."""
     from transformers import AutoConfig
@@ -44,3 +122,5 @@ def register():
         "NeMoSpeechLMForConditionalGeneration",
         f"{_PKG}.model:NeMoSpeechLMForConditionalGeneration",
     )
+
+    _patch_vllm_for_nemo_speechlm_mtp()

nemo/collections/speechlm2/vllm/salm/audio.py

@@ -82,9 +82,12 @@
 
 
 def _ensure_special_tokens(tokenizer: PreTrainedTokenizerBase) -> None:
-    special = [_AUDIO_PLACEHOLDER]
-    existing = set(tokenizer.get_vocab().keys())
-    to_add = [t for t in special if t not in existing]
+    # NOTE: called per request from _call_hf_processor on the API-server event loop.
+    # Use O(1) dict membership; `set(get_vocab().keys())` rebuilt a 131k-entry set
+    # every request (~5-6 ms) purely to check one token. get_vocab() returns vLLM's
+    # cached dict, so membership is O(1).
+    vocab = tokenizer.get_vocab()
+    to_add = [t for t in (_AUDIO_PLACEHOLDER,) if t not in vocab]
     if to_add:
         tokenizer.add_special_tokens({"additional_special_tokens": to_add})
 

nemo/collections/speechlm2/vllm/salm/config.py

@@ -178,6 +178,12 @@ def __init__(
 
         self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
 
+        # vLLM's MTP llm_base_proposer reads image_token_index from the target
+        # model's config to locate multimodal placeholder positions during
+        # speculative decoding. For SpeechLM the <|audio|> token is the first
+        # extra row added above the base backbone vocab.
+        self.image_token_index = self.text_config.vocab_size - _SPEECHLM_EMBED_EXTRA_ROWS
+
     @property
     def llm_architectures(self) -> list[str]:
         """Return the LLM backbone architectures list."""
@@ -186,6 +192,16 @@ def llm_architectures(self) -> list[str]:
     def get_text_config(self, decoder=False) -> PretrainedConfig:
         return self.text_config
 
+    @property
+    def mtp_hybrid_override_pattern(self) -> str:
+        """Hybrid layer pattern for MTP heads, consumed by NemotronHMultiTokenPredictor.
+
+        Reads from the ``mtp.hybrid_override_pattern`` field in config.json.
+        '*' means all-attention; 'M' means all-Mamba2.
+        """
+        mtp_cfg = self.__dict__.get("mtp") or {}
+        return mtp_cfg.get("hybrid_override_pattern", "*") if isinstance(mtp_cfg, dict) else "*"
+
     _ATTR_ALIASES = {
         "rms_norm_eps": "layer_norm_epsilon",
         "layer_norm_eps": "layer_norm_epsilon",

nemo/collections/speechlm2/vllm/salm/model.py

@@ -180,6 +180,33 @@ def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
             return []
         return self._process_audio(audio_input)
 
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Embed token IDs and fuse audio embeddings at placeholder positions.
+
+        Required so that vLLM's MTP speculator probe
+        (``draft_model.embed_input_ids(ids, multimodal_embeddings=None)``)
+        succeeds and ``speculator.supports_mm_inputs`` stays True.
+        Without this method the probe raises AttributeError and the
+        speculator silently falls back to text-only draft mode.
+        """
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+
+        if multimodal_embeddings is None or is_multimodal is None or not is_multimodal.any():
+            return inputs_embeds
+
+        # Concatenate per-audio embedding tensors and overwrite the audio
+        # placeholder token positions with the actual audio embeddings.
+        audio_embeds = torch.cat(list(multimodal_embeddings), dim=0)
+        inputs_embeds = inputs_embeds.clone()
+        inputs_embeds[is_multimodal] = audio_embeds.to(inputs_embeds.dtype)
+        return inputs_embeds
+
     # ── forward / logits ──
 
     def forward(
@@ -222,6 +249,8 @@ def _split_perception_llm(
                 continue
             if name.startswith("perception."):
                 perception[name[len("perception.") :]] = tensor
+            elif name.startswith("llm.mtp.") or name.startswith("mtp."):
+                pass  # MTP draft-head weights; loaded by the speculative draft model, not here
             else:
                 llm.append((name, tensor))
         return perception, llm

nemo/collections/speechlm2/vllm/salm/mtp.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MTP speculative-decoding draft model for NeMo SpeechLM checkpoints.
+
+``NeMoSpeechLMMTP`` wraps vLLM's ``NemotronHMTP`` and adapts the
+weight-loading step for the NeMo checkpoint layout where all LLM
+weights (including MTP layers) carry an ``llm.`` prefix:
+
+    NeMo checkpoint            NemotronHMTP expects
+    ──────────────────────     ────────────────────
+    llm.mtp.layers.0.*    →    mtp.layers.0.*
+    llm.lm_head.weight    →    lm_head.weight
+    llm.model.embed_*     →    (embeddings shared from target model)
+
+Only the ``mtp.*`` and ``lm_head.*`` weights are loaded here; the
+embedding table is shared with the target model by vLLM's MTP framework.
+"""
+
+from collections.abc import Iterable
+
+import torch
+from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP
+from vllm.sequence import IntermediateTensors
+
+from nemo.collections.speechlm2.vllm.salm.audio import _pad_to_vocab_size
+
+
+class NeMoSpeechLMMTP(NemotronHMTP):
+    """NemotronH MTP draft model for NeMo SpeechLM checkpoints.
+
+    Extends NemotronHMTP in two ways:
+
+    * ``load_weights`` strips the ``llm.`` prefix from checkpoint names.
+    * ``embed_input_ids`` fuses audio-feature embeddings into the token
+      embeddings at placeholder positions, exactly like the target model.
+      The MTP heads were trained on the same mixed text+audio embedding
+      stream as the backbone, so the draft must see it too. vLLM probes
+      ``draft_model.embed_input_ids(ids, multimodal_embeddings=None)`` at
+      load time (``llm_base_proposer.load_model``); without this method
+      the probe raises AttributeError and speculative decoding silently
+      falls back to text-only draft inputs, which collapses acceptance
+      rates on audio prompts.
+    """
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Embed token IDs and merge audio embeddings at placeholder positions.
+
+        Mirrors ``NeMoSpeechLMForConditionalGeneration.embed_input_ids``. The
+        embedding table itself is shared from the target model by vLLM's MTP
+        framework, so text-token rows are identical to the target's.
+        """
+        inputs_embeds = self.model.get_input_embeddings(input_ids)
+
+        if multimodal_embeddings is None or is_multimodal is None or not is_multimodal.any():
+            return inputs_embeds
+
+        audio_embeds = torch.cat(list(multimodal_embeddings), dim=0)
+        inputs_embeds = inputs_embeds.clone()
+        inputs_embeds[is_multimodal] = audio_embeds.to(inputs_embeds.dtype)
+        return inputs_embeds
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        lm_head_vocab = None
+        for name, module in self.named_modules():
+            if hasattr(module, "org_vocab_size") and "lm_head" in name:
+                lm_head_vocab = module.org_vocab_size
+                break
+
+        def _strip_llm_prefix(items):
+            for name, tensor in items:
+                if name.startswith("llm."):
+                    name = name[len("llm.") :]
+                if name == "lm_head.weight" and lm_head_vocab is not None:
+                    tensor = _pad_to_vocab_size(tensor, lm_head_vocab)
+                yield name, tensor
+
+        return super().load_weights(_strip_llm_prefix(weights))