Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin by Slyne · Pull Request #15906 · NVIDIA-NeMo/Speech

Slyne · 2026-07-11T23:00:55Z

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

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add MTP speculative decoding support for NeMo SpeechLM models served through vLLM.

Collection: SpeechLM (nemo.collections.speechlm2)

Changelog

Register a NeMo SpeechLM MTP draft-model architecture with vLLM.
Route SpeechLM checkpoints with MTP configuration through vLLM speculative decoding.
Support repeated-layer, multi-head MTP checkpoints.
Fuse audio embeddings into placeholder positions for both target and draft models.
Load MTP checkpoint weights while handling SpeechLM prefixes and vocabulary padding.
Reject incompatible multi-head checkpoints that were not trained with a repeated MTP layer.
Make vLLM plugin registration idempotent.
Add unit tests for registration, configuration routing, weight handling, and multimodal MTP embeddings.

Usage

You can potentially add a usage example below

from nemo.collections.speechlm2.vllm.salm import register

register()

And then start vllm server:

vllm serve /path/to/speechlm-checkpoint \
  --trust-remote-code \
  --speculative-config '{"method":"mtp","num_speculative_tokens":4}'

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

55 focused SpeechLM vLLM plugin tests.
vLLM 0.23.0.
A four-head repeated-layer MTP checkpoint.
The complete 83,212-sample Open-ASR leaderboard.
MTP and no-MTP WER remained equivalent within normal run-to-run variation.

copy-pr-bot · 2026-07-11T23:00:59Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

DongjiGao · 2026-07-24T03:37:46Z

+        inputs_embeds = inputs_embeds.clone()
+        inputs_embeds[is_multimodal] = audio_embeds.to(inputs_embeds.dtype)
+        return inputs_embeds
+


Is the target-model override necessary?
NeMoSpeechLMForConditionalGeneration already inherits the default implementation from SupportsMultiModal

Thanks. Removed.

- 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 <noreply@anthropic.com> Signed-off-by: slyne deng <slyned@nvidia.com>

Signed-off-by: SlyneD <slyned@nvidia.com>

…r repeated MTP head; and fix for v0.20.0 weight loading Signed-off-by: SlyneD <slyned@nvidia.com>

Signed-off-by: SlyneD <slyned@nvidia.com>

Signed-off-by: slyne deng <slyned@nvidia.com>

+    @pytest.fixture(autouse=True)
+    def restore_original_override(self):
+        """Keep the process-local fallback hook isolated between tests."""
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        """A fresh spawn import should delegate unrelated configs to vLLM's native hook."""
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        """A corrupted install must fail clearly instead of recursing into our override."""
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module


Slyne requested a review from DongjiGao July 11, 2026 23:01

github-advanced-security Bot found potential problems Jul 11, 2026

View reviewed changes

Comment thread nemo/collections/speechlm2/vllm/salm/mtp.py Fixed

Comment thread tests/collections/speechlm2/test_vllm_plugin.py Fixed

Slyne force-pushed the slyne/mtp-vllm-eval branch 2 times, most recently from f088f08 to 25fce05 Compare July 20, 2026 20:10

DongjiGao reviewed Jul 25, 2026

View reviewed changes

DongjiGao mentioned this pull request Aug 16, 2026

Update MTP inference for SpeechLM vLLM for Nemotron 3.5 Lightning based models #16068

Open

Slyne added 5 commits August 21, 2026 15:27

fix(speechlm2): remove unused MTP plugin symbols

d7ce2a0

Signed-off-by: SlyneD <slyned@nvidia.com>

allow spec tokens to get value beyond the trained number of tokens fo…

4fa283e

…r repeated MTP head; and fix for v0.20.0 weight loading Signed-off-by: SlyneD <slyned@nvidia.com>

fix based on comments

02a972b

Signed-off-by: SlyneD <slyned@nvidia.com>

fix(speechlm2): harden SALM MTP vLLM support

8a4df2d

Signed-off-by: slyne deng <slyned@nvidia.com>

Slyne force-pushed the slyne/mtp-vllm-eval branch from 436f2ea to 8a4df2d Compare August 24, 2026 18:05

Slyne added 2 commits August 24, 2026 15:26

fix(speechlm2): derive audio token index at runtime

b723aa6

Signed-off-by: slyne deng <slyned@nvidia.com>

fix(speechlm2): address final SALM review findings

ad2514c

Signed-off-by: slyne deng <slyned@nvidia.com>

Slyne mentioned this pull request Aug 25, 2026

Add DFlash speculative decoding support to NeMo SpeechLM vLLM plugin #16129

Draft

6 tasks

fix(speechlm2): make MTP override spawn-pickleable

4ad950b

Signed-off-by: slyne deng <slyned@nvidia.com>

github-advanced-security Bot found potential problems Aug 25, 2026

View reviewed changes

Navigation Menu

Uh oh!

Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin - #15906

Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin#15906
Slyne wants to merge 8 commits into
NVIDIA-NeMo:mainfrom
Slyne:slyne/mtp-vllm-eval

Slyne commented Jul 11, 2026

Uh oh!

copy-pr-bot Bot commented Jul 11, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DongjiGao Jul 24, 2026

Uh oh!

Slyne Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Slyne commented Jul 11, 2026

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

copy-pr-bot Bot commented Jul 11, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DongjiGao Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

Slyne Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants