From 7636c9fb36cf59767016f297fae15ad580ba0e67 Mon Sep 17 00:00:00 2001 From: slyne deng Date: Fri, 10 Jul 2026 14:06:48 -0700 Subject: [PATCH] Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin - register the NeMo SpeechLM MTP draft model and route compatible repeated-layer checkpoints through vLLM speculative decoding - fuse audio embeddings at placeholder positions for target and draft models - load MTP weights while handling SpeechLM checkpoint prefixes and vocabulary padding - make plugin registration idempotent and validate repeated-layer multi-head configurations - add focused plugin tests for registration, config routing, embeddings, and MTP properties Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: slyne deng --- .../speechlm2/vllm/salm/__init__.py | 80 ++++++++ nemo/collections/speechlm2/vllm/salm/audio.py | 9 +- .../collections/speechlm2/vllm/salm/config.py | 16 ++ nemo/collections/speechlm2/vllm/salm/model.py | 29 +++ nemo/collections/speechlm2/vllm/salm/mtp.py | 95 +++++++++ .../collections/speechlm2/test_vllm_plugin.py | 184 ++++++++++++++++++ 6 files changed, 410 insertions(+), 3 deletions(-) create mode 100644 nemo/collections/speechlm2/vllm/salm/mtp.py diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py index 177e6918aa91..2fc13476e9c5 100644 --- a/nemo/collections/speechlm2/vllm/salm/__init__.py +++ b/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() diff --git a/nemo/collections/speechlm2/vllm/salm/audio.py b/nemo/collections/speechlm2/vllm/salm/audio.py index 7ce90e79a2f1..d55f3b827308 100644 --- a/nemo/collections/speechlm2/vllm/salm/audio.py +++ b/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}) diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py index 6d9f55d1b3fc..6cd3d4863a8c 100644 --- a/nemo/collections/speechlm2/vllm/salm/config.py +++ b/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", diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py index a7160f5f8a32..6b1723b88360 100644 --- a/nemo/collections/speechlm2/vllm/salm/model.py +++ b/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 diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py new file mode 100644 index 000000000000..9f76ea2d1cda --- /dev/null +++ b/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)) diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py index 565ee5f08935..c6f51a1e2484 100644 --- a/tests/collections/speechlm2/test_vllm_plugin.py +++ b/tests/collections/speechlm2/test_vllm_plugin.py @@ -591,6 +591,190 @@ def test_register_does_not_load_backbone_config(self, monkeypatch): from_pretrained.assert_not_called() +@pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed") +class TestMTPPlugin: + """Tests for NeMo SpeechLM MTP speculative-decoding support.""" + + class _HFConfigLike: + """Minimal stand-in for a HuggingFace PretrainedConfig.""" + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + def update(self, d): + for k, v in d.items(): + setattr(self, k, v) + + def test_mtp_patch_registers_model(self, monkeypatch): + """register() should add NeMoSpeechLMMTPModel to the model registry.""" + from transformers import AutoConfig + from vllm.model_executor.models.registry import ModelRegistry + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + + register() + + assert "NeMoSpeechLMMTPModel" in ModelRegistry.get_supported_archs() + + def test_mtp_patch_extends_mtp_model_types(self, monkeypatch): + """register() should add 'nemo_speechlm_mtp' to vLLM's MTPModelTypes Literal.""" + from typing import get_args + + import vllm.config.speculative as _spec_mod + from transformers import AutoConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + + register() + + assert "nemo_speechlm_mtp" in get_args(_spec_mod.MTPModelTypes) + + def test_patched_override_routes_nemo_mtp_config(self, monkeypatch): + """hf_config_override should rewrite nemo_speechlm configs with MTP heads.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={"num_nextn_predict_layers": 1, "use_repeated_layer": True}, + ) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.model_type == "nemo_speechlm_mtp" + assert result.architectures == ["NeMoSpeechLMMTPModel"] + assert result.n_predict == 1 + assert result.num_nextn_predict_layers == 1 + + def test_patched_override_no_mtp_falls_through(self, monkeypatch): + """hf_config_override should not alter non-MTP configs.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + original_calls = [] + _orig = SpeculativeConfig.hf_config_override + + def _recording_orig(cfg): + original_calls.append(cfg) + return cfg + + monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig)) + register() + + hf_cfg = self._HFConfigLike(model_type="nemo_speechlm", mtp={"num_nextn_predict_layers": 0}) + SpeculativeConfig.hf_config_override(hf_cfg) + + assert len(original_calls) == 1 + + def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeypatch): + """hf_config_override should raise for multi-head checkpoints without use_repeated_layer.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={"num_nextn_predict_layers": 3, "use_repeated_layer": False}, + ) + with pytest.raises(ValueError, match="use_repeated_layer"): + SpeculativeConfig.hf_config_override(hf_cfg) + + def test_mtp_override_registration_is_idempotent(self, monkeypatch): + """register() should not repeatedly wrap the config override.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError())) + register() + first_override = SpeculativeConfig.hf_config_override + register() + + assert SpeculativeConfig.hf_config_override is first_override + + def test_embed_input_ids_text_only(self): + """embed_input_ids with no audio embeddings should return plain text embeddings.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + m = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.arange(6, dtype=torch.float).reshape(3, 2) + m.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone()) + + result = m.embed_input_ids(torch.tensor([1, 2, 3]), multimodal_embeddings=None) + + assert result.shape == (3, 2) + assert torch.equal(result, base_embeds) + + def test_embed_input_ids_fuses_audio(self): + """embed_input_ids should replace placeholder positions with audio embeddings.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + m = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.zeros(4, 2) + m.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone()) + + audio_feat = torch.ones(2, 2) * 9.0 + is_audio = torch.tensor([False, True, True, False]) + result = m.embed_input_ids( + torch.tensor([0, 1, 2, 3]), + multimodal_embeddings=[audio_feat], + is_multimodal=is_audio, + ) + + assert torch.equal(result[0], torch.zeros(2)) + assert torch.equal(result[1], torch.ones(2) * 9.0) + assert torch.equal(result[2], torch.ones(2) * 9.0) + assert torch.equal(result[3], torch.zeros(2)) + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_mtp_hybrid_override_pattern_from_config(self): + """mtp_hybrid_override_pattern should read hybrid_override_pattern from mtp config dict.""" + cfg = NeMoSpeechLMConfig( + **_DEFAULT_CONFIG_KWARGS, + mtp={"num_nextn_predict_layers": 1, "hybrid_override_pattern": "M*M"}, + ) + assert cfg.mtp_hybrid_override_pattern == "M*M" + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_mtp_hybrid_override_pattern_default_all_attention(self): + """mtp_hybrid_override_pattern should default to '*' (all-attention) when absent.""" + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) + assert cfg.mtp_hybrid_override_pattern == "*" + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_image_token_index_is_base_vocab_size(self): + """image_token_index should equal the backbone base vocab size (before padding).""" + import importlib + + config_mod = importlib.import_module("nemo.collections.speechlm2.vllm.salm.config") + extra_rows = config_mod._SPEECHLM_EMBED_EXTRA_ROWS + + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) + base_vocab = cfg.text_config.vocab_size - extra_rows + assert cfg.image_token_index == base_vocab + + class _FakeTokenizer: def __init__(self): self.added_special_tokens = None