From 9cc7dca53378d4ee64f5e6887bd4db2303f25994 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Fri, 10 Jul 2026 14:06:48 -0700
Subject: [PATCH 01/20] 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 <noreply@anthropic.com>

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 .../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 3c65829d2dd5..e86ea35b3a0c 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -705,6 +705,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

From d7ce2a0bf0754405f85f185e688bdad2a769e5ca Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Mon, 20 Jul 2026 12:23:17 -0700
Subject: [PATCH 02/20] fix(speechlm2): remove unused MTP plugin symbols

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 nemo/collections/speechlm2/vllm/salm/mtp.py     | 1 -
 tests/collections/speechlm2/test_vllm_plugin.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py
index 9f76ea2d1cda..3301e66912a8 100644
--- a/nemo/collections/speechlm2/vllm/salm/mtp.py
+++ b/nemo/collections/speechlm2/vllm/salm/mtp.py
@@ -32,7 +32,6 @@
 
 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
 
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index e86ea35b3a0c..ab743bfd8312 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -778,7 +778,6 @@ def test_patched_override_no_mtp_falls_through(self, monkeypatch):
 
         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)

From 4fa283ed35a8b980fab67d7a4f8d270d61823621 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Mon, 27 Jul 2026 15:37:52 -0700
Subject: [PATCH 03/20] allow spec tokens to get value beyond the trained
 number of tokens for repeated MTP head; and fix for v0.20.0 weight loading

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 .../speechlm2/vllm/salm/__init__.py           |  8 +--
 nemo/collections/speechlm2/vllm/salm/mtp.py   | 40 +++++++++-----
 .../collections/speechlm2/test_vllm_plugin.py | 52 +++++++++++++++++++
 3 files changed, 83 insertions(+), 17 deletions(-)

diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index 2fc13476e9c5..f2493ee8bdda 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -76,9 +76,11 @@ def _patched_override(hf_config):
                     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,
+                            # Size of the physical MTP block that vLLM reuses. A
+                            # repeated-layer checkpoint ships one shared head even
+                            # when it was trained for multiple next-token positions,
+                            # so arbitrary inference K values must be multiples of 1.
+                            "n_predict": 1 if use_repeated_layer else 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
diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py
index 3301e66912a8..70d13167b500 100644
--- a/nemo/collections/speechlm2/vllm/salm/mtp.py
+++ b/nemo/collections/speechlm2/vllm/salm/mtp.py
@@ -22,10 +22,10 @@
     ──────────────────────     ────────────────────
     llm.mtp.layers.0.*    →    mtp.layers.0.*
     llm.lm_head.weight    →    lm_head.weight
-    llm.model.embed_*     →    (embeddings shared from target model)
+    llm.model.embed_*     →    backbone.embeddings.*
 
-Only the ``mtp.*`` and ``lm_head.*`` weights are loaded here; the
-embedding table is shared with the target model by vLLM's MTP framework.
+The embedding alias is required by vLLM's Nemotron-H MTP loader even though
+the proposer subsequently shares the table with the target model.
 """
 
 from collections.abc import Iterable
@@ -36,6 +36,26 @@
 from nemo.collections.speechlm2.vllm.salm.audio import _pad_to_vocab_size
 
 
+def _remap_nemo_mtp_weights(
+    items: Iterable[tuple[str, torch.Tensor]], target_vocab: int | None = None
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Map exported NeMo SpeechLM names to ``NemotronHMTP`` aliases."""
+    for name, tensor in items:
+        if name.startswith("llm."):
+            name = name[len("llm.") :]
+
+        # NemotronHMTP.load_weights only admits embedding names containing
+        # ``embeddings`` and then maps this backbone alias to
+        # ``model.embed_tokens``. Passing model.embed_tokens through directly
+        # is silently skipped and DefaultModelLoader reports it uninitialized.
+        if name == "model.embed_tokens.weight":
+            name = "backbone.embeddings.weight"
+
+        if name in {"backbone.embeddings.weight", "lm_head.weight"} and target_vocab is not None:
+            tensor = _pad_to_vocab_size(tensor, target_vocab)
+        yield name, tensor
+
+
 class NeMoSpeechLMMTP(NemotronHMTP):
     """NemotronH MTP draft model for NeMo SpeechLM checkpoints.
 
@@ -77,18 +97,10 @@ def embed_input_ids(
         return inputs_embeds
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        lm_head_vocab = None
+        target_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
+                target_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))
+        return super().load_weights(_remap_nemo_mtp_weights(weights, target_vocab))
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index ab743bfd8312..a8c70333244a 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -769,6 +769,25 @@ def test_patched_override_routes_nemo_mtp_config(self, monkeypatch):
         assert result.n_predict == 1
         assert result.num_nextn_predict_layers == 1
 
+    def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeypatch):
+        """Repeated-layer training depth must not constrain inference-time K."""
+        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": 4, "use_repeated_layer": True},
+        )
+        result = SpeculativeConfig.hf_config_override(hf_cfg)
+
+        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
@@ -860,6 +879,39 @@ def test_embed_input_ids_fuses_audio(self):
         assert torch.equal(result[2], torch.ones(2) * 9.0)
         assert torch.equal(result[3], torch.zeros(2))
 
+    def test_mtp_weight_remap_uses_vllm_embedding_alias(self):
+        """Exported SpeechLM embeddings must pass NemotronHMTP's name filter."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights
+
+        tensor = torch.ones(2, 3)
+        remapped = dict(
+            _remap_nemo_mtp_weights(
+                [
+                    ("llm.model.embed_tokens.weight", tensor),
+                    ("llm.mtp.layers.0.enorm.weight", tensor),
+                    ("llm.lm_head.weight", tensor),
+                ]
+            )
+        )
+
+        assert set(remapped) == {
+            "backbone.embeddings.weight",
+            "mtp.layers.0.enorm.weight",
+            "lm_head.weight",
+        }
+        assert remapped["backbone.embeddings.weight"] is tensor
+
+        padded = dict(
+            _remap_nemo_mtp_weights(
+                [("llm.model.embed_tokens.weight", tensor), ("llm.lm_head.weight", tensor)],
+                target_vocab=5,
+            )
+        )
+        assert padded["backbone.embeddings.weight"].shape == (5, 3)
+        assert padded["lm_head.weight"].shape == (5, 3)
+
     @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."""

From 02a972bcccc70bed05e54f174520f18fdab6187a Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Mon, 27 Jul 2026 16:59:45 -0700
Subject: [PATCH 04/20] fix based on comments

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 .../speechlm2/vllm/salm/__init__.py           | 12 +++++----
 .../collections/speechlm2/vllm/salm/config.py |  5 ++--
 nemo/collections/speechlm2/vllm/salm/model.py | 27 -------------------
 nemo/collections/speechlm2/vllm/salm/mtp.py   |  7 ++---
 4 files changed, 13 insertions(+), 38 deletions(-)

diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index f2493ee8bdda..148c9274daa8 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -49,20 +49,22 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     import vllm.config.speculative as _spec_mod
     from vllm.config.speculative import SpeculativeConfig
 
-    # 1 ── Extend MTPModelTypes -------------------------------------------
+    # Extend vLLM's recognized MTP model types.
     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 ----------------------------------------
+    # Route SpeechLM MTP checkpoints through SpeculativeConfig.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
+                mtp_cfg = getattr(hf_config, "mtp", None)
+                if not isinstance(mtp_cfg, dict):
+                    mtp_cfg = {}
+                n_predict = mtp_cfg.get("num_nextn_predict_layers", 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:
@@ -97,7 +99,7 @@ def _patched_override(hf_config):
         _patched_override._nemo_speechlm_mtp_override = True
         SpeculativeConfig.hf_config_override = staticmethod(_patched_override)
 
-    # 3 ── Register NeMoSpeechLMMTPModel ----------------------------------
+    # Register the SpeechLM MTP draft architecture with vLLM.
     from vllm.model_executor.models.registry import ModelRegistry
 
     ModelRegistry.register_model(
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index 6cd3d4863a8c..362e41122b5a 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -176,13 +176,12 @@ def __init__(
             if num_layers > 0:
                 self.text_config.layer_types = ["attention"] * num_layers
 
-        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
+        self.image_token_index = self.text_config.vocab_size
+        self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
 
     @property
     def llm_architectures(self) -> list[str]:
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index 6b1723b88360..aabe50c8b206 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -180,33 +180,6 @@ 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(
diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py
index 70d13167b500..9331f5016441 100644
--- a/nemo/collections/speechlm2/vllm/salm/mtp.py
+++ b/nemo/collections/speechlm2/vllm/salm/mtp.py
@@ -82,9 +82,10 @@ def embed_input_ids(
     ) -> 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.
+        The target model inherits this fusion from ``SupportsMultiModal``; the
+        draft must implement it itself because ``NemotronHMTP`` is not
+        multimodal. The embedding table is shared from the target by vLLM's
+        MTP framework, so text-token rows stay identical.
         """
         inputs_embeds = self.model.get_input_embeddings(input_ids)
 

From 8a4df2d1a4a1bb8ca8153f96b250214526628b28 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 10:15:43 -0700
Subject: [PATCH 05/20] fix(speechlm2): harden SALM MTP vLLM support

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   |  49 ++-
 .../speechlm2/vllm/salm/__init__.py           |  41 +-
 .../collections/speechlm2/vllm/salm/config.py |  48 ++-
 nemo/collections/speechlm2/vllm/salm/model.py |   2 +-
 nemo/collections/speechlm2/vllm/salm/mtp.py   | 117 +++++-
 tests/collections/speechlm2/test_to_hf.py     | 126 +++++-
 .../collections/speechlm2/test_vllm_plugin.py | 394 +++++++++++++++++-
 7 files changed, 722 insertions(+), 55 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index e0734fbce70e..279221e16dfd 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -102,6 +102,18 @@ def _canonical_torch_dtype_name(dtype: str | torch.dtype) -> str:
 def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[str, Any]:
     """Build the exported root config without mutating the training config."""
     config = OmegaConf.to_container(model.cfg) if isinstance(model.cfg, DictConfig) else deepcopy(model.cfg)
+    mtp_cfg = config.get("mtp")
+    llm_config = getattr(getattr(model, "llm", None), "config", None)
+    actual_mtp_pattern = getattr(llm_config, "mtp_hybrid_override_pattern", None)
+    if isinstance(mtp_cfg, dict) and mtp_cfg.get("enabled", False) and actual_mtp_pattern is not None:
+        if not isinstance(actual_mtp_pattern, str) or not actual_mtp_pattern:
+            raise ValueError(
+                f"Built LLM has invalid mtp_hybrid_override_pattern={actual_mtp_pattern!r}; cannot export it."
+            )
+        # A preserved checkpoint-native MTP head can differ from the recipe's
+        # requested replacement pattern. Persist the pattern of the head that
+        # was actually built so vLLM instantiates the matching physical layers.
+        mtp_cfg["hybrid_override_pattern"] = actual_mtp_pattern
     dtype_name = _canonical_torch_dtype_name(dtype)
     config["dtype"] = dtype_name
     config["torch_dtype"] = dtype_name
@@ -135,13 +147,15 @@ def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) ->
     llm_config.save_pretrained(str(llm_backbone_dir))
 
 
-def _detect_vllm_architecture(model_cfg: dict) -> str:
-    """Determine the vLLM plugin model class for the checkpoint.
+def _inspect_vllm_backbone(model_cfg: dict) -> tuple[str, int]:
+    """Determine the vLLM plugin class and model embedding vocabulary size.
 
     The SALM plugin registers a single architecture name and selects between
     transformer and hybrid backends at instantiation time, so this function
-    just verifies the backbone config is reachable and returns the unified
-    name; the hybrid-vs-transformer split is handled inside the plugin.
+    verifies the backbone config is reachable and returns the unified name
+    plus the model config's vocabulary size. The latter is authoritative for
+    the embedding-table bound; tokenizer ``vocab_size`` can exclude added
+    tokens and need not equal it.
 
     Raises:
         ValueError: if the HF config can't be loaded or has no 'architectures'.
@@ -160,8 +174,11 @@ def _detect_vllm_architecture(model_cfg: dict) -> str:
     archs = getattr(llm_cfg, "architectures", [])
     if not archs:
         raise ValueError(f"HF config for {pretrained_llm!r} has empty 'architectures'.")
+    vocab_size = getattr(llm_cfg, "vocab_size", None)
+    if isinstance(vocab_size, bool) or not isinstance(vocab_size, int) or vocab_size <= 0:
+        raise ValueError(f"HF config for {pretrained_llm!r} has invalid 'vocab_size': {vocab_size!r}.")
 
-    return "NeMoSpeechLMForConditionalGeneration"
+    return "NeMoSpeechLMForConditionalGeneration", vocab_size
 
 
 def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
@@ -179,6 +196,7 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     """
     from transformers import AutoTokenizer
 
+    from nemo.collections.speechlm2.vllm.salm.config import _SPEECHLM_EMBED_EXTRA_ROWS
     from nemo.utils import logging as LOG
 
     output_dir = Path(output_dir)
@@ -197,13 +215,12 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
     if (llm_backbone_dir / "config.json").exists():
         arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
-    arch = _detect_vllm_architecture(arch_model_cfg)
+    arch, base_vocab_size = _inspect_vllm_backbone(arch_model_cfg)
     config_path = output_dir / "config.json"
     config = json.loads(config_path.read_text())
     config["model_type"] = "nemo_speechlm"
     config["architectures"] = [arch]
     config["audio_locator_tag"] = audio_token
-    config_path.write_text(json.dumps(config, indent=2) + "\n")
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [
@@ -213,9 +230,25 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     ]
     if existing:
         LOG.info("Overwriting existing files in %s: %s", output_dir, existing)
-    tok = AutoTokenizer.from_pretrained(pretrained_llm, trust_remote_code=True)
+    tokenizer_src = model_cfg.get("tokenizer_path") or pretrained_llm
+    tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True)
     if audio_token not in tok.get_vocab():
         tok.add_special_tokens({"additional_special_tokens": [audio_token]})
+    audio_token_id = tok.get_vocab().get(audio_token)
+    if isinstance(audio_token_id, bool) or not isinstance(audio_token_id, int) or audio_token_id < 0:
+        raise ValueError(f"Tokenizer did not assign a valid ID to audio token {audio_token!r}.")
+    padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
+    if audio_token_id >= padded_vocab_size:
+        raise ValueError(
+            f"Audio token ID {audio_token_id} is outside the SpeechLM embedding table with "
+            f"{padded_vocab_size} rows. Reduce the tokenizer's added-token count before training/export."
+        )
+    # Persist the exact placeholder ID under the SALM-specific name. The
+    # runtime config exposes vLLM's historical ``image_token_index`` spelling
+    # as a compatibility property for its EAGLE-style proposer.
+    config.pop("image_token_index", None)
+    config["audio_token_index"] = audio_token_id
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
     tok.save_pretrained(str(output_dir))
     # Newer transformers splits long chat_template into a separate
     # ``chat_template.jinja`` file; inline it back and drop the file.
diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index 148c9274daa8..bbbe64b194f0 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -23,13 +23,17 @@
 Backbone-specific behavior is selected at instantiation time.
 """
 
+import logging
+
 _PKG = "nemo.collections.speechlm2.vllm.salm"
+_LOG = logging.getLogger(__name__)
 
 
 def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
 
-    Three patches are applied:
+    Releases without ``MTPModelTypes`` return without patching so the ordinary
+    SpeechLM target model remains usable. Otherwise, three patches are applied:
 
     1. ``MTPModelTypes`` — the Literal type that guards the MTP detection
        branch in ``SpeculativeConfig.__post_init__`` is extended to include
@@ -49,6 +53,17 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     import vllm.config.speculative as _spec_mod
     from vllm.config.speculative import SpeculativeConfig
 
+    # MTP support was added incrementally across vLLM releases. Keep the
+    # ordinary SpeechLM target-model plugin usable on releases that do not yet
+    # expose this type guard; speculative decoding will remain unavailable and
+    # vLLM will report that if the user tries to enable it.
+    if not hasattr(_spec_mod, "MTPModelTypes"):
+        _LOG.warning(
+            "This vLLM release does not expose MTPModelTypes; NeMo SpeechLM "
+            "will be registered without MTP speculative-decoding support."
+        )
+        return
+
     # Extend vLLM's recognized MTP model types.
     old_args = get_args(_spec_mod.MTPModelTypes)
     if "nemo_speechlm_mtp" not in old_args:
@@ -64,8 +79,12 @@ def _patched_override(hf_config):
                 mtp_cfg = getattr(hf_config, "mtp", None)
                 if not isinstance(mtp_cfg, dict):
                     mtp_cfg = {}
-                n_predict = mtp_cfg.get("num_nextn_predict_layers", 0)
-                if n_predict > 0:
+                # Match SALMAutomodel's training defaults exactly: merely
+                # retaining a recipe depth does not enable MTP, while an enabled
+                # block with no explicit depth constructs one logical head.
+                mtp_enabled = bool(mtp_cfg.get("enabled", False))
+                n_predict = mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0)
+                if mtp_enabled and 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(
@@ -82,13 +101,15 @@ def _patched_override(hf_config):
                             # repeated-layer checkpoint ships one shared head even
                             # when it was trained for multiple next-token positions,
                             # so arbitrary inference K values must be multiples of 1.
-                            "n_predict": 1 if use_repeated_layer else 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.
+                            # Consequently vLLM defaults to K=1 when K is omitted;
+                            # callers should set num_speculative_tokens explicitly.
+                            "n_predict": 1,
+                            # Physical MTP prediction steps to instantiate. Repeated-layer checkpoints
+                            # ship one shared step (one mtp.layers.* module per hybrid-pattern character)
+                            # that is reapplied every speculative iteration, exactly as vLLM drives its
+                            # MTP draft. This also shadows the backbone text_config's
+                            # num_nextn_predict_layers (e.g. 4), which would otherwise trip the
+                            # single-step assert in NemotronHMultiTokenPredictor.
                             "num_nextn_predict_layers": 1,
                             "architectures": ["NeMoSpeechLMMTPModel"],
                         }
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index 362e41122b5a..1e35b9d68691 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -40,8 +40,9 @@
 _AUDIO_PLACEHOLDER = "<|audio|>"
 
 # Number of extra embedding rows the SpeechLM adds on top of the backbone's
-# native vocab during training: ``<|audio|>`` locator plus headroom for other
-# special tokens and TensorCore-friendly alignment.
+# native vocab during training for special tokens and alignment. New exports
+# record the exact ``<|audio|>`` token ID; legacy exports placed it in the first
+# extra row and fall back to the base vocabulary size.
 _SPEECHLM_EMBED_EXTRA_ROWS = 10
 
 
@@ -73,6 +74,8 @@ def __init__(
         pretrained_llm: str | None = None,
         pretrained_asr: str | None = None,
         audio_locator_tag: str | None = None,
+        audio_token_index: int | None = None,
+        image_token_index: int | None = None,
         prompt_format: str | None = None,
         pretrained_weights: bool | None = None,
         lora: dict | None = None,
@@ -90,6 +93,8 @@ def __init__(
             perception is None
             and lora is None
             and encoder_chunk_size_seconds is None
+            and audio_token_index is None
+            and image_token_index is None
             and not kwargs
             and all(value is None for value in required_fields.values())
         )
@@ -111,6 +116,7 @@ def __init__(
             self.pretrained_llm = None
             self.pretrained_asr = None
             self.audio_locator_tag = None
+            self.audio_token_index = None
             self.prompt_format = None
             self.pretrained_weights = None
             self.lora = None
@@ -176,12 +182,31 @@ def __init__(
             if num_layers > 0:
                 self.text_config.layer_types = ["attention"] * num_layers
 
-        # 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
-        self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
+        # New exports persist the exact audio placeholder ID under the
+        # modality-correct name. Accept the historical vLLM field for
+        # compatibility, and keep the base-vocab fallback for older NeMo
+        # checkpoints that persisted neither field and appended <|audio|>
+        # immediately after the backbone vocabulary.
+        base_vocab_size = int(self.text_config.vocab_size)
+        if audio_token_index is not None and image_token_index is not None and audio_token_index != image_token_index:
+            raise ValueError(
+                f"audio_token_index={audio_token_index} conflicts with legacy "
+                f"image_token_index={image_token_index}."
+            )
+        if audio_token_index is None:
+            audio_token_index = image_token_index
+        if audio_token_index is None:
+            audio_token_index = base_vocab_size
+        if isinstance(audio_token_index, bool) or not isinstance(audio_token_index, int):
+            raise ValueError(f"audio_token_index must be an integer, got {audio_token_index!r}.")
+        padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
+        if not 0 <= audio_token_index < padded_vocab_size:
+            raise ValueError(
+                f"audio_token_index={audio_token_index} is outside the SpeechLM embedding table "
+                f"with {padded_vocab_size} rows."
+            )
+        self.audio_token_index = audio_token_index
+        self.text_config.vocab_size = padded_vocab_size
 
     @property
     def llm_architectures(self) -> list[str]:
@@ -191,6 +216,11 @@ def llm_architectures(self) -> list[str]:
     def get_text_config(self, decoder=False) -> PretrainedConfig:
         return self.text_config
 
+    @property
+    def image_token_index(self) -> int | None:
+        """Compatibility alias required by vLLM's EAGLE/MTP proposer."""
+        return self.audio_token_index
+
     @property
     def mtp_hybrid_override_pattern(self) -> str:
         """Hybrid layer pattern for MTP heads, consumed by NemotronHMultiTokenPredictor.
@@ -228,6 +258,8 @@ def __getattr__(self, name):
             "pretrained_llm",
             "pretrained_asr",
             "audio_locator_tag",
+            "audio_token_index",
+            "image_token_index",
             "prompt_format",
             "pretrained_weights",
             "text_config",
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index aabe50c8b206..381a70457b66 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -222,7 +222,7 @@ def _split_perception_llm(
                 continue
             if name.startswith("perception."):
                 perception[name[len("perception.") :]] = tensor
-            elif name.startswith("llm.mtp.") or name.startswith("mtp."):
+            elif name.startswith("llm.mtp."):
                 pass  # MTP draft-head weights; loaded by the speculative draft model, not here
             else:
                 llm.append((name, tensor))
diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py
index 9331f5016441..f6a630a67341 100644
--- a/nemo/collections/speechlm2/vllm/salm/mtp.py
+++ b/nemo/collections/speechlm2/vllm/salm/mtp.py
@@ -15,7 +15,7 @@
 """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
+weight-loading step for the NeMo SpeechLM checkpoint layout where all LLM
 weights (including MTP layers) carry an ``llm.`` prefix:
 
     NeMo checkpoint            NemotronHMTP expects
@@ -33,16 +33,77 @@
 import torch
 from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP
 
+try:
+    from vllm.model_executor.models.utils import _merge_multimodal_embeddings as _vllm_merge_multimodal_embeddings
+except ImportError:  # pragma: no cover - exercised by monkeypatching the resolved helper
+    _vllm_merge_multimodal_embeddings = None
+
 from nemo.collections.speechlm2.vllm.salm.audio import _pad_to_vocab_size
 
 
+def _flatten_multimodal_embeddings(embeddings) -> torch.Tensor:
+    """Flatten vLLM-style nested embedding tensors on every dimension but the last."""
+    if isinstance(embeddings, torch.Tensor):
+        return embeddings.flatten(0, -2)
+    return torch.cat(tuple(_flatten_multimodal_embeddings(item) for item in embeddings))
+
+
+def _merge_multimodal_embeddings(
+    inputs_embeds: torch.Tensor,
+    multimodal_embeddings,
+    is_multimodal: torch.Tensor,
+) -> torch.Tensor:
+    """Use vLLM's merge helper when available, with a compatible fallback for older releases."""
+    if _vllm_merge_multimodal_embeddings is not None:
+        return _vllm_merge_multimodal_embeddings(inputs_embeds, multimodal_embeddings, is_multimodal)
+
+    mm_embeds_flat = _flatten_multimodal_embeddings(multimodal_embeddings)
+    try:
+        inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=inputs_embeds.dtype)
+    except RuntimeError as error:
+        actual_tokens = len(mm_embeds_flat)
+        expected_tokens = is_multimodal.sum().item()
+        if actual_tokens != expected_tokens:
+            raise ValueError(
+                f"Attempted to assign {actual_tokens} multimodal tokens to {expected_tokens} placeholders"
+            ) from error
+        raise ValueError("Error during multimodal embedding index assignment") from error
+    return inputs_embeds
+
+
 def _remap_nemo_mtp_weights(
-    items: Iterable[tuple[str, torch.Tensor]], target_vocab: int | None = None
+    items: Iterable[tuple[str, torch.Tensor]],
+    target_vocab: int | None = None,
+    expected_layer_modules: int | None = None,
 ) -> Iterable[tuple[str, torch.Tensor]]:
-    """Map exported NeMo SpeechLM names to ``NemotronHMTP`` aliases."""
+    """Select and map exported SpeechLM draft weights to ``NemotronHMTP`` aliases.
+
+    ``expected_layer_modules`` is the number of sublayers instantiated for one
+    EAGLE prediction step (the hybrid pattern length), not the logical draft K.
+    Rejecting higher checkpoint indices prevents vLLM from silently dropping
+    weights from a distinct multi-head checkpoint routed as a repeated head.
+    Non-``llm.`` tensors (including perception weights) are intentionally
+    omitted because the target model loads them separately.
+    """
     for name, tensor in items:
-        if name.startswith("llm."):
-            name = name[len("llm.") :]
+        if not name.startswith("llm."):
+            continue
+        name = name[len("llm.") :]
+
+        parts = name.split(".")
+        if expected_layer_modules is not None and len(parts) > 2 and parts[:2] == ["mtp", "layers"]:
+            try:
+                layer_index = int(parts[2])
+            except ValueError:
+                layer_index = None
+            if layer_index is not None and layer_index >= expected_layer_modules:
+                raise ValueError(
+                    f"Checkpoint contains {name!r}, but the configured repeated MTP step instantiates "
+                    f"only {expected_layer_modules} hybrid-pattern layer module(s). This looks like a "
+                    f"distinct multi-head checkpoint or stale hybrid-pattern metadata. Verify that the "
+                    f"exported mtp.hybrid_override_pattern describes the built head; vLLM's NemotronH "
+                    f"MTP proposer can reuse only one EAGLE-style prediction step."
+                )
 
         # NemotronHMTP.load_weights only admits embedding names containing
         # ``embeddings`` and then maps this backbone alias to
@@ -61,7 +122,7 @@ class NeMoSpeechLMMTP(NemotronHMTP):
 
     Extends NemotronHMTP in two ways:
 
-    * ``load_weights`` strips the ``llm.`` prefix from checkpoint names.
+    * ``load_weights`` strips the NeMo SpeechLM ``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
@@ -89,19 +150,39 @@ def embed_input_ids(
         """
         inputs_embeds = self.model.get_input_embeddings(input_ids)
 
-        if multimodal_embeddings is None or is_multimodal is None or not is_multimodal.any():
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
             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
+        if is_multimodal is None:
+            raise ValueError("is_multimodal is required when multimodal_embeddings are provided.")
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        target_vocab = None
-        for name, module in self.named_modules():
-            if hasattr(module, "org_vocab_size") and "lm_head" in name:
-                target_vocab = module.org_vocab_size
-                break
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
 
-        return super().load_weights(_remap_nemo_mtp_weights(weights, target_vocab))
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load only the SALM-prefixed weights required by the reusable draft head."""
+        # NeMoSpeechLMConfig delegates this to the already padded text-config
+        # vocabulary used to construct NemotronHMTP. Reading config directly
+        # avoids coupling padding to vLLM's internal module names.
+        target_vocab = int(self.config.vocab_size)
+        if target_vocab <= 0:
+            raise ValueError(f"Draft model vocabulary size must be positive, got {target_vocab}.")
+
+        # vLLM instantiates one physical module per character in the configured
+        # hybrid pattern. Derive the checkpoint bound from the serialized
+        # configuration instead of silently depending on a predictor-internal
+        # attribute that could drift between vLLM releases.
+        pattern = self.config.mtp_hybrid_override_pattern
+        if not isinstance(pattern, str) or not pattern:
+            raise ValueError(f"mtp_hybrid_override_pattern must be a non-empty string, got {pattern!r}.")
+        expected_layer_modules = len(pattern)
+        return super().load_weights(
+            _remap_nemo_mtp_weights(
+                weights,
+                target_vocab,
+                expected_layer_modules=expected_layer_modules,
+            )
+        )
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index 8cd3b08bd7cb..c3fe9545644a 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -14,12 +14,13 @@
 """Unit tests for ``examples/speechlm2/to_hf.py::prepare_for_vllm``.
 
 The script lives under ``examples/`` (not an importable package), so we load
-it via ``importlib`` and patch ``AutoTokenizer`` / ``_detect_vllm_architecture``
+it via ``importlib`` and patch ``AutoTokenizer`` / ``_inspect_vllm_backbone``
 to avoid any network or real-model dependencies.
 """
 import importlib.util
 import json
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -54,8 +55,10 @@ def __init__(
         split_chat_template=False,
         tokenizer_class="Qwen2Tokenizer",
         eos_token_id=42,
+        base_vocab_size=None,
     ):
         self._vocab = {tok: i for i, tok in enumerate(vocab_tokens)}
+        self.vocab_size = len(self._vocab) if base_vocab_size is None else base_vocab_size
         self._chat_template = chat_template
         self._split_chat_template = split_chat_template
         self._tokenizer_class = tokenizer_class
@@ -130,6 +133,30 @@ class _FakeExportModel:
     llm = type("_FakeLLM", (), {"config": _FakeLLMConfig()})()
 
 
+def test_hf_export_config_persists_built_mtp_pattern_without_mutating_recipe():
+    """A preserved native head's physical pattern must override stale recipe metadata."""
+    model = SimpleNamespace(
+        cfg={"mtp": {"enabled": True, "hybrid_override_pattern": "*"}},
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E")),
+    )
+
+    config = to_hf._hf_export_config(model, "bfloat16")
+
+    assert config["mtp"]["hybrid_override_pattern"] == "*E"
+    assert model.cfg["mtp"]["hybrid_override_pattern"] == "*"
+
+
+@pytest.mark.parametrize("pattern", ["", 3])
+def test_hf_export_config_rejects_invalid_built_mtp_pattern(pattern):
+    model = SimpleNamespace(
+        cfg={"mtp": {"enabled": True, "hybrid_override_pattern": "*"}},
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern=pattern)),
+    )
+
+    with pytest.raises(ValueError, match="mtp_hybrid_override_pattern"):
+        to_hf._hf_export_config(model, "bfloat16")
+
+
 def test_save_hf_checkpoint_writes_llm_backbone_config(tmp_path):
     cfg = to_hf.HfExportConfig(
         class_path="fake.Class",
@@ -185,15 +212,41 @@ def test_prepare_for_vllm_missing_audio_locator_tag(tmp_path):
         to_hf.prepare_for_vllm(str(tmp_path), {"pretrained_llm": "fake-model"})
 
 
+def test_inspect_vllm_backbone_returns_model_embedding_vocab():
+    """Exporter bounds must use the model config, not tokenizer.vocab_size."""
+    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=151936)
+    with patch("transformers.AutoConfig.from_pretrained", return_value=backbone):
+        architecture, vocab_size = to_hf._inspect_vllm_backbone({"pretrained_llm": "fake-model"})
+
+    assert architecture == "NeMoSpeechLMForConditionalGeneration"
+    assert vocab_size == 151936
+
+
+@pytest.mark.parametrize("vocab_size", [None, True, 0, -1])
+def test_inspect_vllm_backbone_rejects_invalid_model_vocab(vocab_size):
+    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=vocab_size)
+    with (
+        patch("transformers.AutoConfig.from_pretrained", return_value=backbone),
+        pytest.raises(ValueError, match="vocab_size"),
+    ):
+        to_hf._inspect_vllm_backbone({"pretrained_llm": "fake-model"})
+
+
 # ──────────────────────────────────────────────────────────────────────
-# Happy paths (mock AutoTokenizer + _detect_vllm_architecture)
+# Happy paths (mock AutoTokenizer + _inspect_vllm_backbone)
 # ──────────────────────────────────────────────────────────────────────
 
 
-def _run_prepare(tmp_path, fake_tok, arch="NeMoSpeechLMForConditionalGeneration", llm_arch="Qwen2ForCausalLM"):
+def _run_prepare(
+    tmp_path,
+    fake_tok,
+    arch="NeMoSpeechLMForConditionalGeneration",
+    llm_arch="Qwen2ForCausalLM",
+    backbone_vocab_size=100,
+):
     output_dir = _seed_output_dir(tmp_path, llm_arch=llm_arch)
     with (
-        patch.object(to_hf, "_detect_vllm_architecture", return_value=arch),
+        patch.object(to_hf, "_inspect_vllm_backbone", return_value=(arch, backbone_vocab_size)),
         patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok),
     ):
         to_hf.prepare_for_vllm(
@@ -204,12 +257,14 @@ def _run_prepare(tmp_path, fake_tok, arch="NeMoSpeechLMForConditionalGeneration"
 
 
 def test_prepare_for_vllm_patches_config_json(tmp_path):
-    """config.json gets model_type, architectures, and audio_locator_tag."""
+    """config.json gets model metadata and the tokenizer's exact audio-token ID."""
     output_dir = _run_prepare(tmp_path, _FakeTokenizer())
     cfg = json.loads((output_dir / "config.json").read_text())
     assert cfg["model_type"] == "nemo_speechlm"
     assert cfg["architectures"] == ["NeMoSpeechLMForConditionalGeneration"]
     assert cfg["audio_locator_tag"] == AUDIO_TOKEN
+    assert cfg["audio_token_index"] == 0
+    assert "image_token_index" not in cfg
     # Original LLM fields are preserved.
     assert cfg["hidden_size"] == 2048
 
@@ -229,6 +284,67 @@ def test_prepare_for_vllm_skips_add_if_audio_token_already_in_vocab(tmp_path):
     assert fake_tok.add_special_tokens_calls == []
 
 
+def test_prepare_for_vllm_records_existing_audio_token_id(tmp_path):
+    """An existing audio token need not sit at the backbone vocabulary boundary."""
+    fake_tok = _FakeTokenizer(vocab_tokens=["<pad>", AUDIO_TOKEN, "<extra>"])
+    output_dir = _run_prepare(tmp_path, fake_tok)
+
+    cfg = json.loads((output_dir / "config.json").read_text())
+    assert cfg["audio_token_index"] == 1
+
+
+def test_prepare_for_vllm_rejects_invalid_audio_token_id(tmp_path):
+    """An invalid tokenizer mapping should fail during export, not at serving time."""
+    fake_tok = _FakeTokenizer(vocab_tokens=[AUDIO_TOKEN])
+    fake_tok._vocab[AUDIO_TOKEN] = -1
+
+    with pytest.raises(ValueError, match="valid ID"):
+        _run_prepare(tmp_path, fake_tok)
+
+
+def test_prepare_for_vllm_rejects_audio_token_outside_padded_embeddings(tmp_path):
+    """The exported placeholder ID must fit the embedding table created by the plugin."""
+    tokens = [f"<extra_{i}>" for i in range(11)] + [AUDIO_TOKEN]
+    fake_tok = _FakeTokenizer(vocab_tokens=tokens, base_vocab_size=1)
+
+    with pytest.raises(ValueError, match="outside the SpeechLM embedding table"):
+        _run_prepare(tmp_path, fake_tok, backbone_vocab_size=1)
+
+
+def test_prepare_for_vllm_uses_model_vocab_bound_not_tokenizer_base_vocab(tmp_path):
+    """Pre-existing added tokens may exceed tokenizer.vocab_size while fitting model embeddings."""
+    tokens = [f"<added_{i}>" for i in range(30)] + [AUDIO_TOKEN]
+    fake_tok = _FakeTokenizer(vocab_tokens=tokens, base_vocab_size=1)
+    output_dir = _run_prepare(tmp_path, fake_tok, backbone_vocab_size=40)
+
+    cfg = json.loads((output_dir / "config.json").read_text())
+    assert cfg["audio_token_index"] == 30
+
+
+def test_prepare_for_vllm_uses_training_tokenizer_path(tmp_path):
+    """Conversion must persist the tokenizer that supplied training token IDs."""
+    output_dir = _seed_output_dir(tmp_path)
+    fake_tok = _FakeTokenizer()
+    with (
+        patch.object(
+            to_hf,
+            "_inspect_vllm_backbone",
+            return_value=("NeMoSpeechLMForConditionalGeneration", 100),
+        ),
+        patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok) as load_tokenizer,
+    ):
+        to_hf.prepare_for_vllm(
+            str(output_dir),
+            {
+                "pretrained_llm": "fake-model",
+                "tokenizer_path": "custom-tokenizer",
+                "audio_locator_tag": AUDIO_TOKEN,
+            },
+        )
+
+    load_tokenizer.assert_called_once_with("custom-tokenizer", trust_remote_code=True)
+
+
 def test_prepare_for_vllm_tokenizer_config_normalized(tmp_path):
     """tokenizer_config.json has dict-form extra_special_tokens + forced tokenizer_class."""
     output_dir = _run_prepare(tmp_path, _FakeTokenizer(tokenizer_class="TokenizersBackend"))
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index a8c70333244a..0945e7dbfbeb 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -20,6 +20,7 @@
 """
 
 import importlib.util
+import logging
 from types import SimpleNamespace
 
 import pytest
@@ -78,11 +79,19 @@ def test_default_construction_for_hf_serialization(self):
         assert cfg.pretrained_llm is None
         assert cfg.pretrained_asr is None
         assert cfg.audio_locator_tag is None
+        assert cfg.audio_token_index is None
+        assert cfg.image_token_index is None
         assert cfg.prompt_format is None
         assert cfg.pretrained_weights is None
         assert cfg.llm_architectures == []
         assert cfg.get_text_config() is cfg.text_config
 
+    @pytest.mark.parametrize("field", ["audio_token_index", "image_token_index"])
+    def test_token_index_alone_is_not_default_construction(self, field):
+        """Checkpoint data must not be silently discarded by HF's no-arg path."""
+        with pytest.raises(ValueError, match="pretrained_llm"):
+            NeMoSpeechLMConfig(**{field: 42})
+
     def test_loads_text_config(self):
         """Config should load a text_config from the pretrained LLM."""
         cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
@@ -709,6 +718,25 @@ def test_register_does_not_load_backbone_config(self, monkeypatch):
 class TestMTPPlugin:
     """Tests for NeMo SpeechLM MTP speculative-decoding support."""
 
+    @pytest.fixture(autouse=True)
+    def mock_backbone_config(self, monkeypatch):
+        """Keep registration tests independent of Hugging Face network access."""
+        if not _HAS_CONFIG:
+            return
+
+        monkeypatch.setattr(
+            _config_module.AutoConfig,
+            "from_pretrained",
+            lambda *args, **kwargs: SimpleNamespace(
+                architectures=["NemotronHybridForCausalLM"],
+                hidden_size=2048,
+                vocab_size=131072,
+                num_hidden_layers=4,
+                num_key_value_heads=2,
+                layer_norm_epsilon=1e-5,
+            ),
+        )
+
     class _HFConfigLike:
         """Minimal stand-in for a HuggingFace PretrainedConfig."""
 
@@ -748,6 +776,31 @@ def test_mtp_patch_extends_mtp_model_types(self, monkeypatch):
 
         assert "nemo_speechlm_mtp" in get_args(_spec_mod.MTPModelTypes)
 
+    def test_mtp_patch_keeps_target_plugin_usable_without_vllm_mtp_guard(self, monkeypatch, caplog):
+        """Older vLLM releases should degrade cleanly instead of breaking registration."""
+        import vllm.config.speculative as _spec_mod
+        from vllm.model_executor.models.registry import ModelRegistry
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        registered_archs = []
+        original_register_model = ModelRegistry.register_model
+
+        def _record_registration(architecture, model):
+            registered_archs.append(architecture)
+            return original_register_model(architecture, model)
+
+        monkeypatch.setattr(ModelRegistry, "register_model", _record_registration)
+        monkeypatch.delattr(_spec_mod, "MTPModelTypes")
+
+        with caplog.at_level(logging.WARNING, logger="nemo.collections.speechlm2.vllm.salm"):
+            register()
+
+        assert "NeMoSpeechLMForConditionalGeneration" in registered_archs
+        assert "NeMoSpeechLMMTPModel" not in registered_archs
+        assert "NeMoSpeechLMForConditionalGeneration" in ModelRegistry.get_supported_archs()
+        assert "without MTP speculative-decoding support" in caplog.text
+
     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
@@ -760,7 +813,7 @@ def test_patched_override_routes_nemo_mtp_config(self, monkeypatch):
 
         hf_cfg = self._HFConfigLike(
             model_type="nemo_speechlm",
-            mtp={"num_nextn_predict_layers": 1, "use_repeated_layer": True},
+            mtp={"enabled": True, "num_nextn_predict_layers": 1, "use_repeated_layer": True},
         )
         result = SpeculativeConfig.hf_config_override(hf_cfg)
 
@@ -769,6 +822,23 @@ def test_patched_override_routes_nemo_mtp_config(self, monkeypatch):
         assert result.n_predict == 1
         assert result.num_nextn_predict_layers == 1
 
+    def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
+        """An enabled training block without an explicit depth constructs one head."""
+        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={"enabled": True})
+        result = SpeculativeConfig.hf_config_override(hf_cfg)
+
+        assert result.model_type == "nemo_speechlm_mtp"
+        assert result.n_predict == 1
+        assert result.num_nextn_predict_layers == 1
+
     def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeypatch):
         """Repeated-layer training depth must not constrain inference-time K."""
         from transformers import AutoConfig
@@ -781,7 +851,7 @@ def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeyp
 
         hf_cfg = self._HFConfigLike(
             model_type="nemo_speechlm",
-            mtp={"num_nextn_predict_layers": 4, "use_repeated_layer": True},
+            mtp={"enabled": True, "num_nextn_predict_layers": 4, "use_repeated_layer": True},
         )
         result = SpeculativeConfig.hf_config_override(hf_cfg)
 
@@ -810,6 +880,58 @@ def _recording_orig(cfg):
 
         assert len(original_calls) == 1
 
+    def test_patched_override_depth_without_enabled_flag_falls_through(self, monkeypatch):
+        """A retained recipe depth must not enable a head that training did not construct."""
+        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 = []
+
+        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": 4, "use_repeated_layer": True},
+        )
+        result = SpeculativeConfig.hf_config_override(hf_cfg)
+
+        assert result.model_type == "nemo_speechlm"
+        assert original_calls == [hf_cfg]
+
+    def test_patched_override_explicitly_disabled_mtp_falls_through(self, monkeypatch):
+        """An exported recipe depth must not override mtp.enabled=false."""
+        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 = []
+
+        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={"enabled": False, "num_nextn_predict_layers": 4, "use_repeated_layer": True},
+        )
+        result = SpeculativeConfig.hf_config_override(hf_cfg)
+
+        assert result.model_type == "nemo_speechlm"
+        assert original_calls == [hf_cfg]
+
     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
@@ -822,7 +944,7 @@ def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeyp
 
         hf_cfg = self._HFConfigLike(
             model_type="nemo_speechlm",
-            mtp={"num_nextn_predict_layers": 3, "use_repeated_layer": False},
+            mtp={"enabled": True, "num_nextn_predict_layers": 3, "use_repeated_layer": False},
         )
         with pytest.raises(ValueError, match="use_repeated_layer"):
             SpeculativeConfig.hf_config_override(hf_cfg)
@@ -879,6 +1001,140 @@ def test_embed_input_ids_fuses_audio(self):
         assert torch.equal(result[2], torch.ones(2) * 9.0)
         assert torch.equal(result[3], torch.zeros(2))
 
+    def test_embed_input_ids_fuses_multiple_audio_chunks(self):
+        """embed_input_ids should preserve vLLM's nested multimodal embedding order."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP
+
+        m = object.__new__(NeMoSpeechLMMTP)
+        base_embeds = torch.zeros(5, 2)
+        m.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone())
+
+        audio_feats = [[torch.ones(1, 2) * 3.0], [torch.ones(2, 2) * 7.0]]
+        is_audio = torch.tensor([True, False, True, True, False])
+        result = m.embed_input_ids(
+            torch.tensor([0, 1, 2, 3, 4]),
+            multimodal_embeddings=audio_feats,
+            is_multimodal=is_audio,
+        )
+
+        assert torch.equal(result[0], torch.ones(2) * 3.0)
+        assert torch.equal(result[1], torch.zeros(2))
+        assert torch.equal(result[2], torch.ones(2) * 7.0)
+        assert torch.equal(result[3], torch.ones(2) * 7.0)
+        assert torch.equal(result[4], torch.zeros(2))
+
+    def test_embed_input_ids_requires_multimodal_mask(self):
+        """Audio embeddings without placeholder positions should fail clearly."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP
+
+        m = object.__new__(NeMoSpeechLMMTP)
+        m.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
+
+        with pytest.raises(ValueError, match="is_multimodal"):
+            m.embed_input_ids(
+                torch.tensor([0, 1]),
+                multimodal_embeddings=[torch.ones(1, 2)],
+            )
+
+    def test_embed_input_ids_ignores_embeddings_without_placeholder_positions(self):
+        """vLLM's merge semantics leave embeddings unchanged for an all-text mask."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP
+
+        model = object.__new__(NeMoSpeechLMMTP)
+        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
+
+        result = model.embed_input_ids(
+            torch.tensor([0, 1]),
+            multimodal_embeddings=[torch.ones(1, 2)],
+            is_multimodal=torch.tensor([False, False]),
+        )
+
+        assert torch.equal(result, torch.zeros(2, 2))
+
+    def test_embed_input_ids_fallback_supports_nested_audio_chunks(self, monkeypatch):
+        """Older vLLM releases without the private merge helper retain nested-input behavior."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
+
+        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
+        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
+        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(4, 2))
+        result = model.embed_input_ids(
+            torch.tensor([0, 1, 2, 3]),
+            multimodal_embeddings=[[torch.ones(1, 2) * 3], [torch.ones(1, 2) * 7]],
+            is_multimodal=torch.tensor([True, False, True, False]),
+        )
+
+        assert torch.equal(result[0], torch.ones(2) * 3)
+        assert torch.equal(result[1], torch.zeros(2))
+        assert torch.equal(result[2], torch.ones(2) * 7)
+        assert torch.equal(result[3], torch.zeros(2))
+
+    def test_embed_input_ids_fallback_preserves_all_text_mask_noop(self, monkeypatch):
+        """The compatibility fallback should match vLLM for an all-False mask."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
+
+        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
+        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
+        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
+        result = model.embed_input_ids(
+            torch.tensor([0, 1]),
+            multimodal_embeddings=[torch.ones(1, 2)],
+            is_multimodal=torch.tensor([False, False]),
+        )
+
+        assert torch.equal(result, torch.zeros(2, 2))
+
+    def test_embed_input_ids_fallback_rejects_embedding_count_mismatch(self, monkeypatch):
+        """The compatibility fallback should report mismatched audio and placeholder counts."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
+
+        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
+        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
+        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
+
+        with pytest.raises(ValueError, match="2 multimodal tokens to 1 placeholders"):
+            model.embed_input_ids(
+                torch.tensor([0, 1]),
+                multimodal_embeddings=[torch.ones(2, 2)],
+                is_multimodal=torch.tensor([True, False]),
+            )
+
+    def test_target_weight_split_excludes_only_salm_mtp_and_extra_state(self):
+        """The target loader should retain LLM weights while routing SALM draft weights away."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        perception_tensor = torch.ones(1)
+        llm_tensor = torch.ones(2)
+        bare_mtp_tensor = torch.ones(3)
+        perception, llm = NeMoSpeechLMForConditionalGeneration._split_perception_llm(
+            [
+                ("perception.encoder.weight", perception_tensor),
+                ("llm.mtp.layers.0.weight", torch.ones(4)),
+                ("llm.model.layers.0.weight", llm_tensor),
+                ("mtp.layers.0.weight", bare_mtp_tensor),
+                ("llm.model.layers.0._extra_state", torch.ones(5)),
+            ]
+        )
+
+        assert perception == {"encoder.weight": perception_tensor}
+        assert [name for name, _ in llm] == ["llm.model.layers.0.weight", "mtp.layers.0.weight"]
+        assert llm[0][1] is llm_tensor
+        assert llm[1][1] is bare_mtp_tensor
+
     def test_mtp_weight_remap_uses_vllm_embedding_alias(self):
         """Exported SpeechLM embeddings must pass NemotronHMTP's name filter."""
         import torch
@@ -912,6 +1168,102 @@ def test_mtp_weight_remap_uses_vllm_embedding_alias(self):
         assert padded["backbone.embeddings.weight"].shape == (5, 3)
         assert padded["lm_head.weight"].shape == (5, 3)
 
+    def test_mtp_weight_remap_rejects_distinct_head_layers(self):
+        """Unsupported distinct heads in a NeMo SpeechLM export should fail loudly."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights
+
+        tensor = torch.ones(2, 3)
+        with pytest.raises(ValueError, match="distinct multi-head"):
+            list(
+                _remap_nemo_mtp_weights(
+                    [("llm.mtp.layers.1.enorm.weight", tensor)],
+                    expected_layer_modules=1,
+                )
+            )
+
+    def test_mtp_weight_remap_ignores_non_salm_names(self):
+        """The draft loader supports final NeMo SpeechLM exports, not bare backbone checkpoints."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights
+
+        assert list(_remap_nemo_mtp_weights([("mtp.layers.0.norm.weight", torch.ones(1))])) == []
+
+    def test_mtp_weight_remap_allows_all_sublayers_in_hybrid_pattern(self):
+        """Hybrid-pattern sublayers belong to one reusable prediction step."""
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights
+
+        tensor = torch.ones(2, 3)
+        remapped = dict(
+            _remap_nemo_mtp_weights(
+                [(f"llm.mtp.layers.{i}.norm.weight", tensor) for i in range(3)],
+                expected_layer_modules=3,
+            )
+        )
+        assert set(remapped) == {f"mtp.layers.{i}.norm.weight" for i in range(3)}
+
+    def test_mtp_load_weights_bounds_layers_from_serialized_pattern(self, monkeypatch):
+        """The real loader wiring should derive its physical-layer bound from config."""
+        import torch
+        from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP
+
+        model = object.__new__(NeMoSpeechLMMTP)
+        object.__setattr__(
+            model,
+            "config",
+            SimpleNamespace(mtp_hybrid_override_pattern="*E", vocab_size=3),
+        )
+        captured = []
+
+        def _capture_weights(self, weights):
+            captured.extend(weights)
+            return set()
+
+        monkeypatch.setattr(NemotronHMTP, "load_weights", _capture_weights)
+        tensor = torch.ones(1, 2)
+        NeMoSpeechLMMTP.load_weights(
+            model,
+            [
+                ("llm.model.embed_tokens.weight", tensor),
+                ("llm.mtp.layers.0.norm.weight", tensor),
+                ("llm.mtp.layers.1.norm.weight", tensor),
+                ("llm.lm_head.weight", tensor),
+            ],
+        )
+        assert [name for name, _ in captured] == [
+            "backbone.embeddings.weight",
+            "mtp.layers.0.norm.weight",
+            "mtp.layers.1.norm.weight",
+            "lm_head.weight",
+        ]
+        captured_tensors = dict(captured)
+        assert captured_tensors["backbone.embeddings.weight"].shape == (3, 2)
+        assert captured_tensors["lm_head.weight"].shape == (3, 2)
+        assert captured_tensors["mtp.layers.0.norm.weight"] is tensor
+        assert captured_tensors["mtp.layers.1.norm.weight"] is tensor
+
+        with pytest.raises(ValueError, match="distinct multi-head"):
+            NeMoSpeechLMMTP.load_weights(model, [("llm.mtp.layers.2.norm.weight", tensor)])
+
+    @pytest.mark.parametrize("pattern", [None, "", 3])
+    def test_mtp_load_weights_rejects_invalid_serialized_pattern(self, monkeypatch, pattern):
+        from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP
+
+        from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP
+
+        model = object.__new__(NeMoSpeechLMMTP)
+        object.__setattr__(model, "config", SimpleNamespace(mtp_hybrid_override_pattern=pattern, vocab_size=3))
+        monkeypatch.setattr(NemotronHMTP, "load_weights", lambda self, weights: set())
+
+        with pytest.raises(ValueError, match="non-empty string"):
+            NeMoSpeechLMMTP.load_weights(model, [])
+
     @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."""
@@ -928,8 +1280,8 @@ def test_mtp_hybrid_override_pattern_default_all_attention(self):
         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)."""
+    def test_audio_token_index_legacy_fallback_is_base_vocab_size(self):
+        """Legacy exports should fall back to the backbone base vocab size."""
         import importlib
 
         config_mod = importlib.import_module("nemo.collections.speechlm2.vllm.salm.config")
@@ -937,8 +1289,40 @@ def test_image_token_index_is_base_vocab_size(self):
 
         cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
         base_vocab = cfg.text_config.vocab_size - extra_rows
+        assert cfg.audio_token_index == base_vocab
         assert cfg.image_token_index == base_vocab
 
+    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+    def test_audio_token_index_uses_exported_tokenizer_id(self):
+        """New exports should use the tokenizer's exact placeholder ID."""
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, audio_token_index=42)
+        assert cfg.audio_token_index == 42
+        assert cfg.image_token_index == 42
+        assert cfg.to_dict()["audio_token_index"] == 42
+        assert "image_token_index" not in cfg.to_dict()
+
+    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+    def test_legacy_image_token_index_is_accepted_as_vllm_alias(self):
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=42)
+        assert cfg.audio_token_index == 42
+        assert cfg.image_token_index == 42
+
+    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+    @pytest.mark.parametrize("field", ["audio_token_index", "image_token_index"])
+    @pytest.mark.parametrize("token_index", [True, "42", -1, 131082])
+    def test_audio_token_index_rejects_invalid_type_or_range(self, field, token_index):
+        with pytest.raises(ValueError, match="audio_token_index"):
+            NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, **{field: token_index})
+
+    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+    def test_conflicting_audio_and_legacy_image_token_indices_raise(self):
+        with pytest.raises(ValueError, match="conflicts"):
+            NeMoSpeechLMConfig(
+                **_DEFAULT_CONFIG_KWARGS,
+                audio_token_index=42,
+                image_token_index=43,
+            )
+
 
 class _FakeTokenizer:
     def __init__(self):

From b723aa6bfa80a49ef31c1174cfc3cd7e594156b6 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 15:26:12 -0700
Subject: [PATCH 06/20] fix(speechlm2): derive audio token index at runtime

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   | 36 +++------
 .../collections/speechlm2/vllm/salm/config.py | 59 ++++++--------
 tests/collections/speechlm2/test_to_hf.py     | 76 ++-----------------
 .../collections/speechlm2/test_vllm_plugin.py | 47 ++----------
 4 files changed, 44 insertions(+), 174 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index 279221e16dfd..b7546604ac2e 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -147,15 +147,13 @@ def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) ->
     llm_config.save_pretrained(str(llm_backbone_dir))
 
 
-def _inspect_vllm_backbone(model_cfg: dict) -> tuple[str, int]:
-    """Determine the vLLM plugin class and model embedding vocabulary size.
+def _detect_vllm_architecture(model_cfg: dict) -> str:
+    """Determine the vLLM plugin model class for the checkpoint.
 
     The SALM plugin registers a single architecture name and selects between
     transformer and hybrid backends at instantiation time, so this function
-    verifies the backbone config is reachable and returns the unified name
-    plus the model config's vocabulary size. The latter is authoritative for
-    the embedding-table bound; tokenizer ``vocab_size`` can exclude added
-    tokens and need not equal it.
+    just verifies the backbone config is reachable and returns the unified
+    name; the hybrid-vs-transformer split is handled inside the plugin.
 
     Raises:
         ValueError: if the HF config can't be loaded or has no 'architectures'.
@@ -174,11 +172,8 @@ def _inspect_vllm_backbone(model_cfg: dict) -> tuple[str, int]:
     archs = getattr(llm_cfg, "architectures", [])
     if not archs:
         raise ValueError(f"HF config for {pretrained_llm!r} has empty 'architectures'.")
-    vocab_size = getattr(llm_cfg, "vocab_size", None)
-    if isinstance(vocab_size, bool) or not isinstance(vocab_size, int) or vocab_size <= 0:
-        raise ValueError(f"HF config for {pretrained_llm!r} has invalid 'vocab_size': {vocab_size!r}.")
 
-    return "NeMoSpeechLMForConditionalGeneration", vocab_size
+    return "NeMoSpeechLMForConditionalGeneration"
 
 
 def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
@@ -196,7 +191,6 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     """
     from transformers import AutoTokenizer
 
-    from nemo.collections.speechlm2.vllm.salm.config import _SPEECHLM_EMBED_EXTRA_ROWS
     from nemo.utils import logging as LOG
 
     output_dir = Path(output_dir)
@@ -215,12 +209,15 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
     if (llm_backbone_dir / "config.json").exists():
         arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
-    arch, base_vocab_size = _inspect_vllm_backbone(arch_model_cfg)
+    arch = _detect_vllm_architecture(arch_model_cfg)
     config_path = output_dir / "config.json"
     config = json.loads(config_path.read_text())
     config["model_type"] = "nemo_speechlm"
     config["architectures"] = [arch]
     config["audio_locator_tag"] = audio_token
+    config.pop("audio_token_index", None)
+    config.pop("image_token_index", None)
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [
@@ -234,21 +231,6 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True)
     if audio_token not in tok.get_vocab():
         tok.add_special_tokens({"additional_special_tokens": [audio_token]})
-    audio_token_id = tok.get_vocab().get(audio_token)
-    if isinstance(audio_token_id, bool) or not isinstance(audio_token_id, int) or audio_token_id < 0:
-        raise ValueError(f"Tokenizer did not assign a valid ID to audio token {audio_token!r}.")
-    padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
-    if audio_token_id >= padded_vocab_size:
-        raise ValueError(
-            f"Audio token ID {audio_token_id} is outside the SpeechLM embedding table with "
-            f"{padded_vocab_size} rows. Reduce the tokenizer's added-token count before training/export."
-        )
-    # Persist the exact placeholder ID under the SALM-specific name. The
-    # runtime config exposes vLLM's historical ``image_token_index`` spelling
-    # as a compatibility property for its EAGLE-style proposer.
-    config.pop("image_token_index", None)
-    config["audio_token_index"] = audio_token_id
-    config_path.write_text(json.dumps(config, indent=2) + "\n")
     tok.save_pretrained(str(output_dir))
     # Newer transformers splits long chat_template into a separate
     # ``chat_template.jinja`` file; inline it back and drop the file.
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index 1e35b9d68691..e4353e56b5a8 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -39,10 +39,10 @@
 # silently rendering the wrong placeholder at request time.
 _AUDIO_PLACEHOLDER = "<|audio|>"
 
-# Number of extra embedding rows the SpeechLM adds on top of the backbone's
-# native vocab during training for special tokens and alignment. New exports
-# record the exact ``<|audio|>`` token ID; legacy exports placed it in the first
-# extra row and fall back to the base vocabulary size.
+# Historical serving-time headroom above the backbone vocabulary. vLLM builds
+# the target and draft embedding tables at this padded size, and the weight
+# loader zero-pads the smaller training tensors to match. ``<|audio|>`` is the
+# first padded ID; the remaining rows are unused.
 _SPEECHLM_EMBED_EXTRA_ROWS = 10
 
 
@@ -74,8 +74,6 @@ def __init__(
         pretrained_llm: str | None = None,
         pretrained_asr: str | None = None,
         audio_locator_tag: str | None = None,
-        audio_token_index: int | None = None,
-        image_token_index: int | None = None,
         prompt_format: str | None = None,
         pretrained_weights: bool | None = None,
         lora: dict | None = None,
@@ -93,8 +91,6 @@ def __init__(
             perception is None
             and lora is None
             and encoder_chunk_size_seconds is None
-            and audio_token_index is None
-            and image_token_index is None
             and not kwargs
             and all(value is None for value in required_fields.values())
         )
@@ -116,7 +112,6 @@ def __init__(
             self.pretrained_llm = None
             self.pretrained_asr = None
             self.audio_locator_tag = None
-            self.audio_token_index = None
             self.prompt_format = None
             self.pretrained_weights = None
             self.lora = None
@@ -182,31 +177,7 @@ def __init__(
             if num_layers > 0:
                 self.text_config.layer_types = ["attention"] * num_layers
 
-        # New exports persist the exact audio placeholder ID under the
-        # modality-correct name. Accept the historical vLLM field for
-        # compatibility, and keep the base-vocab fallback for older NeMo
-        # checkpoints that persisted neither field and appended <|audio|>
-        # immediately after the backbone vocabulary.
-        base_vocab_size = int(self.text_config.vocab_size)
-        if audio_token_index is not None and image_token_index is not None and audio_token_index != image_token_index:
-            raise ValueError(
-                f"audio_token_index={audio_token_index} conflicts with legacy "
-                f"image_token_index={image_token_index}."
-            )
-        if audio_token_index is None:
-            audio_token_index = image_token_index
-        if audio_token_index is None:
-            audio_token_index = base_vocab_size
-        if isinstance(audio_token_index, bool) or not isinstance(audio_token_index, int):
-            raise ValueError(f"audio_token_index must be an integer, got {audio_token_index!r}.")
-        padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
-        if not 0 <= audio_token_index < padded_vocab_size:
-            raise ValueError(
-                f"audio_token_index={audio_token_index} is outside the SpeechLM embedding table "
-                f"with {padded_vocab_size} rows."
-            )
-        self.audio_token_index = audio_token_index
-        self.text_config.vocab_size = padded_vocab_size
+        self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
 
     @property
     def llm_architectures(self) -> list[str]:
@@ -218,8 +189,23 @@ def get_text_config(self, decoder=False) -> PretrainedConfig:
 
     @property
     def image_token_index(self) -> int | None:
-        """Compatibility alias required by vLLM's EAGLE/MTP proposer."""
-        return self.audio_token_index
+        """Return the audio placeholder ID expected by vLLM 0.26's MTP proposer.
+
+        SpeechLM training appends ``<|audio|>`` directly after the backbone
+        vocabulary. The vision-era vLLM proposer calls this field
+        ``image_token_index`` even for an audio multimodal target.
+        """
+        vocab_size = getattr(self.text_config, "vocab_size", None)
+        if vocab_size is None:
+            return None
+        return int(vocab_size) - _SPEECHLM_EMBED_EXTRA_ROWS
+
+    @image_token_index.setter
+    def image_token_index(self, value: int | None) -> None:
+        """Accept vLLM's runtime target-to-draft copy without serializing it."""
+        expected = self.image_token_index
+        if expected is not None and value != expected:
+            raise ValueError(f"image_token_index={value!r} does not match the backbone vocabulary size {expected}.")
 
     @property
     def mtp_hybrid_override_pattern(self) -> str:
@@ -258,7 +244,6 @@ def __getattr__(self, name):
             "pretrained_llm",
             "pretrained_asr",
             "audio_locator_tag",
-            "audio_token_index",
             "image_token_index",
             "prompt_format",
             "pretrained_weights",
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index c3fe9545644a..f2028ab8d86d 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -14,7 +14,7 @@
 """Unit tests for ``examples/speechlm2/to_hf.py::prepare_for_vllm``.
 
 The script lives under ``examples/`` (not an importable package), so we load
-it via ``importlib`` and patch ``AutoTokenizer`` / ``_inspect_vllm_backbone``
+it via ``importlib`` and patch ``AutoTokenizer`` / ``_detect_vllm_architecture``
 to avoid any network or real-model dependencies.
 """
 import importlib.util
@@ -55,10 +55,8 @@ def __init__(
         split_chat_template=False,
         tokenizer_class="Qwen2Tokenizer",
         eos_token_id=42,
-        base_vocab_size=None,
     ):
         self._vocab = {tok: i for i, tok in enumerate(vocab_tokens)}
-        self.vocab_size = len(self._vocab) if base_vocab_size is None else base_vocab_size
         self._chat_template = chat_template
         self._split_chat_template = split_chat_template
         self._tokenizer_class = tokenizer_class
@@ -212,28 +210,8 @@ def test_prepare_for_vllm_missing_audio_locator_tag(tmp_path):
         to_hf.prepare_for_vllm(str(tmp_path), {"pretrained_llm": "fake-model"})
 
 
-def test_inspect_vllm_backbone_returns_model_embedding_vocab():
-    """Exporter bounds must use the model config, not tokenizer.vocab_size."""
-    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=151936)
-    with patch("transformers.AutoConfig.from_pretrained", return_value=backbone):
-        architecture, vocab_size = to_hf._inspect_vllm_backbone({"pretrained_llm": "fake-model"})
-
-    assert architecture == "NeMoSpeechLMForConditionalGeneration"
-    assert vocab_size == 151936
-
-
-@pytest.mark.parametrize("vocab_size", [None, True, 0, -1])
-def test_inspect_vllm_backbone_rejects_invalid_model_vocab(vocab_size):
-    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=vocab_size)
-    with (
-        patch("transformers.AutoConfig.from_pretrained", return_value=backbone),
-        pytest.raises(ValueError, match="vocab_size"),
-    ):
-        to_hf._inspect_vllm_backbone({"pretrained_llm": "fake-model"})
-
-
 # ──────────────────────────────────────────────────────────────────────
-# Happy paths (mock AutoTokenizer + _inspect_vllm_backbone)
+# Happy paths (mock AutoTokenizer + _detect_vllm_architecture)
 # ──────────────────────────────────────────────────────────────────────
 
 
@@ -242,11 +220,10 @@ def _run_prepare(
     fake_tok,
     arch="NeMoSpeechLMForConditionalGeneration",
     llm_arch="Qwen2ForCausalLM",
-    backbone_vocab_size=100,
 ):
     output_dir = _seed_output_dir(tmp_path, llm_arch=llm_arch)
     with (
-        patch.object(to_hf, "_inspect_vllm_backbone", return_value=(arch, backbone_vocab_size)),
+        patch.object(to_hf, "_detect_vllm_architecture", return_value=arch),
         patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok),
     ):
         to_hf.prepare_for_vllm(
@@ -257,13 +234,13 @@ def _run_prepare(
 
 
 def test_prepare_for_vllm_patches_config_json(tmp_path):
-    """config.json gets model metadata and the tokenizer's exact audio-token ID."""
+    """config.json gets model metadata without persisting a token-index field."""
     output_dir = _run_prepare(tmp_path, _FakeTokenizer())
     cfg = json.loads((output_dir / "config.json").read_text())
     assert cfg["model_type"] == "nemo_speechlm"
     assert cfg["architectures"] == ["NeMoSpeechLMForConditionalGeneration"]
     assert cfg["audio_locator_tag"] == AUDIO_TOKEN
-    assert cfg["audio_token_index"] == 0
+    assert "audio_token_index" not in cfg
     assert "image_token_index" not in cfg
     # Original LLM fields are preserved.
     assert cfg["hidden_size"] == 2048
@@ -284,53 +261,12 @@ def test_prepare_for_vllm_skips_add_if_audio_token_already_in_vocab(tmp_path):
     assert fake_tok.add_special_tokens_calls == []
 
 
-def test_prepare_for_vllm_records_existing_audio_token_id(tmp_path):
-    """An existing audio token need not sit at the backbone vocabulary boundary."""
-    fake_tok = _FakeTokenizer(vocab_tokens=["<pad>", AUDIO_TOKEN, "<extra>"])
-    output_dir = _run_prepare(tmp_path, fake_tok)
-
-    cfg = json.loads((output_dir / "config.json").read_text())
-    assert cfg["audio_token_index"] == 1
-
-
-def test_prepare_for_vllm_rejects_invalid_audio_token_id(tmp_path):
-    """An invalid tokenizer mapping should fail during export, not at serving time."""
-    fake_tok = _FakeTokenizer(vocab_tokens=[AUDIO_TOKEN])
-    fake_tok._vocab[AUDIO_TOKEN] = -1
-
-    with pytest.raises(ValueError, match="valid ID"):
-        _run_prepare(tmp_path, fake_tok)
-
-
-def test_prepare_for_vllm_rejects_audio_token_outside_padded_embeddings(tmp_path):
-    """The exported placeholder ID must fit the embedding table created by the plugin."""
-    tokens = [f"<extra_{i}>" for i in range(11)] + [AUDIO_TOKEN]
-    fake_tok = _FakeTokenizer(vocab_tokens=tokens, base_vocab_size=1)
-
-    with pytest.raises(ValueError, match="outside the SpeechLM embedding table"):
-        _run_prepare(tmp_path, fake_tok, backbone_vocab_size=1)
-
-
-def test_prepare_for_vllm_uses_model_vocab_bound_not_tokenizer_base_vocab(tmp_path):
-    """Pre-existing added tokens may exceed tokenizer.vocab_size while fitting model embeddings."""
-    tokens = [f"<added_{i}>" for i in range(30)] + [AUDIO_TOKEN]
-    fake_tok = _FakeTokenizer(vocab_tokens=tokens, base_vocab_size=1)
-    output_dir = _run_prepare(tmp_path, fake_tok, backbone_vocab_size=40)
-
-    cfg = json.loads((output_dir / "config.json").read_text())
-    assert cfg["audio_token_index"] == 30
-
-
 def test_prepare_for_vllm_uses_training_tokenizer_path(tmp_path):
     """Conversion must persist the tokenizer that supplied training token IDs."""
     output_dir = _seed_output_dir(tmp_path)
     fake_tok = _FakeTokenizer()
     with (
-        patch.object(
-            to_hf,
-            "_inspect_vllm_backbone",
-            return_value=("NeMoSpeechLMForConditionalGeneration", 100),
-        ),
+        patch.object(to_hf, "_detect_vllm_architecture", return_value="NeMoSpeechLMForConditionalGeneration"),
         patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok) as load_tokenizer,
     ):
         to_hf.prepare_for_vllm(
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 0945e7dbfbeb..8bbc95aa3d3c 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -79,19 +79,12 @@ def test_default_construction_for_hf_serialization(self):
         assert cfg.pretrained_llm is None
         assert cfg.pretrained_asr is None
         assert cfg.audio_locator_tag is None
-        assert cfg.audio_token_index is None
         assert cfg.image_token_index is None
         assert cfg.prompt_format is None
         assert cfg.pretrained_weights is None
         assert cfg.llm_architectures == []
         assert cfg.get_text_config() is cfg.text_config
 
-    @pytest.mark.parametrize("field", ["audio_token_index", "image_token_index"])
-    def test_token_index_alone_is_not_default_construction(self, field):
-        """Checkpoint data must not be silently discarded by HF's no-arg path."""
-        with pytest.raises(ValueError, match="pretrained_llm"):
-            NeMoSpeechLMConfig(**{field: 42})
-
     def test_loads_text_config(self):
         """Config should load a text_config from the pretrained LLM."""
         cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
@@ -1280,8 +1273,8 @@ def test_mtp_hybrid_override_pattern_default_all_attention(self):
         assert cfg.mtp_hybrid_override_pattern == "*"
 
     @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
-    def test_audio_token_index_legacy_fallback_is_base_vocab_size(self):
-        """Legacy exports should fall back to the backbone base vocab size."""
+    def test_image_token_index_is_unserialized_backbone_vocab_boundary(self):
+        """vLLM's compatibility property should identify the first padded row."""
         import importlib
 
         config_mod = importlib.import_module("nemo.collections.speechlm2.vllm.salm.config")
@@ -1289,40 +1282,14 @@ def test_audio_token_index_legacy_fallback_is_base_vocab_size(self):
 
         cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
         base_vocab = cfg.text_config.vocab_size - extra_rows
-        assert cfg.audio_token_index == base_vocab
         assert cfg.image_token_index == base_vocab
-
-    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
-    def test_audio_token_index_uses_exported_tokenizer_id(self):
-        """New exports should use the tokenizer's exact placeholder ID."""
-        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, audio_token_index=42)
-        assert cfg.audio_token_index == 42
-        assert cfg.image_token_index == 42
-        assert cfg.to_dict()["audio_token_index"] == 42
+        # vLLM 0.26 copies the target value onto the draft config. The setter
+        # accepts that runtime-only assignment without adding serialized state.
+        cfg.image_token_index = base_vocab
+        assert cfg.image_token_index == base_vocab
+        assert "audio_token_index" not in cfg.to_dict()
         assert "image_token_index" not in cfg.to_dict()
 
-    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
-    def test_legacy_image_token_index_is_accepted_as_vllm_alias(self):
-        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=42)
-        assert cfg.audio_token_index == 42
-        assert cfg.image_token_index == 42
-
-    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
-    @pytest.mark.parametrize("field", ["audio_token_index", "image_token_index"])
-    @pytest.mark.parametrize("token_index", [True, "42", -1, 131082])
-    def test_audio_token_index_rejects_invalid_type_or_range(self, field, token_index):
-        with pytest.raises(ValueError, match="audio_token_index"):
-            NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, **{field: token_index})
-
-    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
-    def test_conflicting_audio_and_legacy_image_token_indices_raise(self):
-        with pytest.raises(ValueError, match="conflicts"):
-            NeMoSpeechLMConfig(
-                **_DEFAULT_CONFIG_KWARGS,
-                audio_token_index=42,
-                image_token_index=43,
-            )
-
 
 class _FakeTokenizer:
     def __init__(self):

From ad2514c0dd50bd057f9326ebd63521fce8e5d4ed Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 16:49:05 -0700
Subject: [PATCH 07/20] fix(speechlm2): address final SALM review findings

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   |  53 ++++++--
 .../speechlm2/vllm/salm/__init__.py           |  17 +--
 .../collections/speechlm2/vllm/salm/config.py |  33 +++--
 nemo/collections/speechlm2/vllm/salm/model.py |   5 +
 nemo/collections/speechlm2/vllm/salm/mtp.py   |  36 +----
 tests/collections/speechlm2/test_to_hf.py     | 125 +++++++++++++++++-
 .../collections/speechlm2/test_vllm_plugin.py | 101 +++-----------
 7 files changed, 212 insertions(+), 158 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index b7546604ac2e..a4ecdb156422 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -114,6 +114,22 @@ def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[
         # requested replacement pattern. Persist the pattern of the head that
         # was actually built so vLLM instantiates the matching physical layers.
         mtp_cfg["hybrid_override_pattern"] = actual_mtp_pattern
+        actual_mtp_depth = getattr(llm_config, "num_nextn_predict_layers", None)
+        if actual_mtp_depth is not None:
+            if isinstance(actual_mtp_depth, bool) or not isinstance(actual_mtp_depth, int) or actual_mtp_depth <= 0:
+                raise ValueError(
+                    f"Built LLM has invalid num_nextn_predict_layers={actual_mtp_depth!r}; cannot export it."
+                )
+            if mtp_cfg.get("use_repeated_layer", False):
+                if actual_mtp_depth != 1:
+                    raise ValueError(
+                        "A repeated MTP head must serialize exactly one physical layer, but the built LLM "
+                        f"declares num_nextn_predict_layers={actual_mtp_depth}."
+                    )
+            else:
+                # For a preserved native head, the recipe depth is advisory.
+                # Export the physical/logical depth that is actually present.
+                mtp_cfg["num_nextn_predict_layers"] = actual_mtp_depth
     dtype_name = _canonical_torch_dtype_name(dtype)
     config["dtype"] = dtype_name
     config["torch_dtype"] = dtype_name
@@ -128,9 +144,8 @@ def save_hf_checkpoint(model: torch.nn.Module, state_dict: dict, cfg: HfExportCo
     target_dtype = str_to_dtype(cfg.dtype)
     state_dict = {k: v.to(target_dtype) for k, v in state_dict.items()}
 
-    save_file(state_dict, output_dir / "model.safetensors")
-
     config = _hf_export_config(model, cfg.dtype)
+    save_file(state_dict, output_dir / "model.safetensors")
     with open(output_dir / "config.json", "w") as f:
         json.dump(config, f, indent=2)
     save_llm_backbone_config(model, output_dir)
@@ -147,16 +162,18 @@ def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) ->
     llm_config.save_pretrained(str(llm_backbone_dir))
 
 
-def _detect_vllm_architecture(model_cfg: dict) -> str:
-    """Determine the vLLM plugin model class for the checkpoint.
+def _detect_vllm_architecture(model_cfg: dict) -> tuple[str, int]:
+    """Determine the vLLM plugin model class and backbone vocabulary size.
 
     The SALM plugin registers a single architecture name and selects between
     transformer and hybrid backends at instantiation time, so this function
-    just verifies the backbone config is reachable and returns the unified
-    name; the hybrid-vs-transformer split is handled inside the plugin.
+    verifies the backbone config is reachable and returns the unified name
+    plus the embedding-table vocabulary bound. The hybrid-vs-transformer split
+    is handled inside the plugin.
 
     Raises:
-        ValueError: if the HF config can't be loaded or has no 'architectures'.
+        ValueError: If the HF config cannot be loaded, has no architecture, or
+            declares an invalid vocabulary size.
     """
     pretrained_llm = model_cfg.get("pretrained_llm", "")
     try:
@@ -172,8 +189,11 @@ def _detect_vllm_architecture(model_cfg: dict) -> str:
     archs = getattr(llm_cfg, "architectures", [])
     if not archs:
         raise ValueError(f"HF config for {pretrained_llm!r} has empty 'architectures'.")
+    vocab_size = getattr(llm_cfg, "vocab_size", None)
+    if isinstance(vocab_size, bool) or not isinstance(vocab_size, int) or vocab_size <= 0:
+        raise ValueError(f"HF config for {pretrained_llm!r} has invalid 'vocab_size': {vocab_size!r}.")
 
-    return "NeMoSpeechLMForConditionalGeneration"
+    return "NeMoSpeechLMForConditionalGeneration", vocab_size
 
 
 def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
@@ -187,10 +207,12 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
         model_cfg: Model config dict (from experiment YAML).
 
     Raises:
-        ValueError: If ``pretrained_llm`` or ``audio_locator_tag`` is missing.
+        ValueError: If required model metadata is missing, or the tokenizer's
+            audio token does not fit the SpeechLM embedding table.
     """
     from transformers import AutoTokenizer
 
+    from nemo.collections.speechlm2.vllm.salm.config import _SPEECHLM_EMBED_EXTRA_ROWS
     from nemo.utils import logging as LOG
 
     output_dir = Path(output_dir)
@@ -209,7 +231,7 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
     if (llm_backbone_dir / "config.json").exists():
         arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
-    arch = _detect_vllm_architecture(arch_model_cfg)
+    arch, base_vocab_size = _detect_vllm_architecture(arch_model_cfg)
     config_path = output_dir / "config.json"
     config = json.loads(config_path.read_text())
     config["model_type"] = "nemo_speechlm"
@@ -217,7 +239,6 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     config["audio_locator_tag"] = audio_token
     config.pop("audio_token_index", None)
     config.pop("image_token_index", None)
-    config_path.write_text(json.dumps(config, indent=2) + "\n")
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [
@@ -231,6 +252,16 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True)
     if audio_token not in tok.get_vocab():
         tok.add_special_tokens({"additional_special_tokens": [audio_token]})
+    audio_token_id = tok.get_vocab().get(audio_token)
+    if isinstance(audio_token_id, bool) or not isinstance(audio_token_id, int) or audio_token_id < 0:
+        raise ValueError(f"Tokenizer did not assign a valid ID to audio token {audio_token!r}.")
+    padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
+    if audio_token_id >= padded_vocab_size:
+        raise ValueError(
+            f"Audio token ID {audio_token_id} is outside the SpeechLM embedding table with "
+            f"{padded_vocab_size} rows. Reduce the tokenizer's added-token count before training/export."
+        )
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
     tok.save_pretrained(str(output_dir))
     # Newer transformers splits long chat_template into a separate
     # ``chat_template.jinja`` file; inline it back and drop the file.
diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index bbbe64b194f0..4ba88d843e4d 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -23,17 +23,13 @@
 Backbone-specific behavior is selected at instantiation time.
 """
 
-import logging
-
 _PKG = "nemo.collections.speechlm2.vllm.salm"
-_LOG = logging.getLogger(__name__)
 
 
 def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
 
-    Releases without ``MTPModelTypes`` return without patching so the ordinary
-    SpeechLM target model remains usable. Otherwise, three patches are applied:
+    Three patches are applied on the supported vLLM 0.19+ releases:
 
     1. ``MTPModelTypes`` — the Literal type that guards the MTP detection
        branch in ``SpeculativeConfig.__post_init__`` is extended to include
@@ -53,17 +49,6 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     import vllm.config.speculative as _spec_mod
     from vllm.config.speculative import SpeculativeConfig
 
-    # MTP support was added incrementally across vLLM releases. Keep the
-    # ordinary SpeechLM target-model plugin usable on releases that do not yet
-    # expose this type guard; speculative decoding will remain unavailable and
-    # vLLM will report that if the user tries to enable it.
-    if not hasattr(_spec_mod, "MTPModelTypes"):
-        _LOG.warning(
-            "This vLLM release does not expose MTPModelTypes; NeMo SpeechLM "
-            "will be registered without MTP speculative-decoding support."
-        )
-        return
-
     # Extend vLLM's recognized MTP model types.
     old_args = get_args(_spec_mod.MTPModelTypes)
     if "nemo_speechlm_mtp" not in old_args:
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index e4353e56b5a8..46bcde9df9e2 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -41,8 +41,8 @@
 
 # Historical serving-time headroom above the backbone vocabulary. vLLM builds
 # the target and draft embedding tables at this padded size, and the weight
-# loader zero-pads the smaller training tensors to match. ``<|audio|>`` is the
-# first padded ID; the remaining rows are unused.
+# loader zero-pads the smaller training tensors to match. The tokenizer's
+# actual audio-token ID is validated against this bound during export.
 _SPEECHLM_EMBED_EXTRA_ROWS = 10
 
 
@@ -101,6 +101,7 @@ def __init__(
         # path; real checkpoint loads replace it below after field validation.
         self.text_config = PretrainedConfig()
         self.is_hybrid = False
+        self._pending_image_token_index = None
 
         super().__init__(**kwargs)
 
@@ -116,6 +117,7 @@ def __init__(
             self.pretrained_weights = None
             self.lora = None
             self.encoder_chunk_size_seconds = None
+            self.__dict__.pop("_pending_image_token_index", None)
             return
 
         for name, value in required_fields.items():
@@ -178,6 +180,13 @@ def __init__(
                 self.text_config.layer_types = ["attention"] * num_layers
 
         self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
+        pending_image_token_index = self.__dict__.pop("_pending_image_token_index", None)
+        if pending_image_token_index is not None and pending_image_token_index != self.image_token_index:
+            raise ValueError(
+                f"image_token_index={pending_image_token_index!r} does not match the backbone vocabulary "
+                f"boundary {self.image_token_index}. Remove this legacy serialized field; SpeechLM derives "
+                f"the vLLM compatibility value at runtime."
+            )
 
     @property
     def llm_architectures(self) -> list[str]:
@@ -189,11 +198,11 @@ def get_text_config(self, decoder=False) -> PretrainedConfig:
 
     @property
     def image_token_index(self) -> int | None:
-        """Return the audio placeholder ID expected by vLLM 0.26's MTP proposer.
+        """Return the vocabulary-boundary value expected by vLLM's MTP proposer.
 
-        SpeechLM training appends ``<|audio|>`` directly after the backbone
-        vocabulary. The vision-era vLLM proposer calls this field
-        ``image_token_index`` even for an audio multimodal target.
+        vLLM calls this compatibility field ``image_token_index`` even for an
+        audio multimodal target. Actual audio locations come from vLLM's
+        placeholder ranges; no token-index field is serialized by SpeechLM.
         """
         vocab_size = getattr(self.text_config, "vocab_size", None)
         if vocab_size is None:
@@ -204,8 +213,16 @@ def image_token_index(self) -> int | None:
     def image_token_index(self, value: int | None) -> None:
         """Accept vLLM's runtime target-to-draft copy without serializing it."""
         expected = self.image_token_index
-        if expected is not None and value != expected:
-            raise ValueError(f"image_token_index={value!r} does not match the backbone vocabulary size {expected}.")
+        if expected is None:
+            # Transformers applies unknown config kwargs in its base-class
+            # constructor, before this wrapper has loaded the real backbone.
+            # Defer validation and discard the temporary value afterwards so
+            # it never becomes serialized state.
+            self._pending_image_token_index = value
+        elif value is not None and value != expected:
+            raise ValueError(
+                f"image_token_index={value!r} does not match the backbone vocabulary boundary {expected}."
+            )
 
     @property
     def mtp_hybrid_override_pattern(self) -> str:
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index 381a70457b66..baa97a2b1dde 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -224,6 +224,11 @@ def _split_perception_llm(
                 perception[name[len("perception.") :]] = tensor
             elif name.startswith("llm.mtp."):
                 pass  # MTP draft-head weights; loaded by the speculative draft model, not here
+            elif name.startswith("mtp."):
+                raise ValueError(
+                    f"Unsupported bare MTP tensor {name!r}; NeMo SpeechLM exports must store draft weights "
+                    f"under the 'llm.mtp.*' namespace."
+                )
             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
index f6a630a67341..2d9ab84db1b3 100644
--- a/nemo/collections/speechlm2/vllm/salm/mtp.py
+++ b/nemo/collections/speechlm2/vllm/salm/mtp.py
@@ -32,45 +32,11 @@
 
 import torch
 from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP
-
-try:
-    from vllm.model_executor.models.utils import _merge_multimodal_embeddings as _vllm_merge_multimodal_embeddings
-except ImportError:  # pragma: no cover - exercised by monkeypatching the resolved helper
-    _vllm_merge_multimodal_embeddings = None
+from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 
 from nemo.collections.speechlm2.vllm.salm.audio import _pad_to_vocab_size
 
 
-def _flatten_multimodal_embeddings(embeddings) -> torch.Tensor:
-    """Flatten vLLM-style nested embedding tensors on every dimension but the last."""
-    if isinstance(embeddings, torch.Tensor):
-        return embeddings.flatten(0, -2)
-    return torch.cat(tuple(_flatten_multimodal_embeddings(item) for item in embeddings))
-
-
-def _merge_multimodal_embeddings(
-    inputs_embeds: torch.Tensor,
-    multimodal_embeddings,
-    is_multimodal: torch.Tensor,
-) -> torch.Tensor:
-    """Use vLLM's merge helper when available, with a compatible fallback for older releases."""
-    if _vllm_merge_multimodal_embeddings is not None:
-        return _vllm_merge_multimodal_embeddings(inputs_embeds, multimodal_embeddings, is_multimodal)
-
-    mm_embeds_flat = _flatten_multimodal_embeddings(multimodal_embeddings)
-    try:
-        inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=inputs_embeds.dtype)
-    except RuntimeError as error:
-        actual_tokens = len(mm_embeds_flat)
-        expected_tokens = is_multimodal.sum().item()
-        if actual_tokens != expected_tokens:
-            raise ValueError(
-                f"Attempted to assign {actual_tokens} multimodal tokens to {expected_tokens} placeholders"
-            ) from error
-        raise ValueError("Error during multimodal embedding index assignment") from error
-    return inputs_embeds
-
-
 def _remap_nemo_mtp_weights(
     items: Iterable[tuple[str, torch.Tensor]],
     target_vocab: int | None = None,
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index f2028ab8d86d..45ad325abb9d 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -98,6 +98,8 @@ def _seed_output_dir(tmp_path, llm_arch="Qwen2ForCausalLM"):
                 "architectures": [llm_arch],
                 "hidden_size": 2048,
                 "num_hidden_layers": 24,
+                "audio_token_index": 17,
+                "image_token_index": 18,
             }
         )
     )
@@ -144,6 +146,55 @@ def test_hf_export_config_persists_built_mtp_pattern_without_mutating_recipe():
     assert model.cfg["mtp"]["hybrid_override_pattern"] == "*"
 
 
+def test_hf_export_config_persists_built_non_repeated_mtp_depth():
+    """A preserved native head's actual depth must override stale recipe metadata."""
+    model = SimpleNamespace(
+        cfg={
+            "mtp": {
+                "enabled": True,
+                "hybrid_override_pattern": "*",
+                "num_nextn_predict_layers": 4,
+                "use_repeated_layer": False,
+            }
+        },
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E", num_nextn_predict_layers=1)),
+    )
+
+    config = to_hf._hf_export_config(model, "bfloat16")
+
+    assert config["mtp"]["num_nextn_predict_layers"] == 1
+    assert model.cfg["mtp"]["num_nextn_predict_layers"] == 4
+
+
+def test_hf_export_config_keeps_logical_depth_for_repeated_mtp():
+    model = SimpleNamespace(
+        cfg={
+            "mtp": {
+                "enabled": True,
+                "hybrid_override_pattern": "*",
+                "num_nextn_predict_layers": 4,
+                "use_repeated_layer": True,
+            }
+        },
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E", num_nextn_predict_layers=1)),
+    )
+
+    config = to_hf._hf_export_config(model, "bfloat16")
+
+    assert config["mtp"]["num_nextn_predict_layers"] == 4
+
+
+@pytest.mark.parametrize("depth", [False, 0, -1])
+def test_hf_export_config_rejects_invalid_built_mtp_depth(depth):
+    model = SimpleNamespace(
+        cfg={"mtp": {"enabled": True, "hybrid_override_pattern": "*"}},
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*", num_nextn_predict_layers=depth)),
+    )
+
+    with pytest.raises(ValueError, match="num_nextn_predict_layers"):
+        to_hf._hf_export_config(model, "bfloat16")
+
+
 @pytest.mark.parametrize("pattern", ["", 3])
 def test_hf_export_config_rejects_invalid_built_mtp_pattern(pattern):
     model = SimpleNamespace(
@@ -155,6 +206,24 @@ def test_hf_export_config_rejects_invalid_built_mtp_pattern(pattern):
         to_hf._hf_export_config(model, "bfloat16")
 
 
+def test_save_hf_checkpoint_validates_config_before_writing_weights(tmp_path):
+    model = SimpleNamespace(
+        cfg={"mtp": {"enabled": True, "hybrid_override_pattern": "*"}},
+        llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="")),
+    )
+    cfg = to_hf.HfExportConfig(
+        class_path="fake.Class",
+        ckpt_path="fake.ckpt",
+        ckpt_config="fake.yaml",
+        output_dir=str(tmp_path),
+    )
+
+    with pytest.raises(ValueError, match="mtp_hybrid_override_pattern"):
+        to_hf.save_hf_checkpoint(model, {"weight": torch.zeros(1)}, cfg)
+
+    assert not (tmp_path / "model.safetensors").exists()
+
+
 def test_save_hf_checkpoint_writes_llm_backbone_config(tmp_path):
     cfg = to_hf.HfExportConfig(
         class_path="fake.Class",
@@ -210,6 +279,26 @@ def test_prepare_for_vllm_missing_audio_locator_tag(tmp_path):
         to_hf.prepare_for_vllm(str(tmp_path), {"pretrained_llm": "fake-model"})
 
 
+def test_detect_vllm_architecture_returns_model_embedding_vocab():
+    """Exporter bounds must use the model config, not tokenizer.vocab_size."""
+    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=151936)
+    with patch("transformers.AutoConfig.from_pretrained", return_value=backbone):
+        architecture, vocab_size = to_hf._detect_vllm_architecture({"pretrained_llm": "fake-model"})
+
+    assert architecture == "NeMoSpeechLMForConditionalGeneration"
+    assert vocab_size == 151936
+
+
+@pytest.mark.parametrize("vocab_size", [None, True, 0, -1])
+def test_detect_vllm_architecture_rejects_invalid_model_vocab(vocab_size):
+    backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=vocab_size)
+    with (
+        patch("transformers.AutoConfig.from_pretrained", return_value=backbone),
+        pytest.raises(ValueError, match="vocab_size"),
+    ):
+        to_hf._detect_vllm_architecture({"pretrained_llm": "fake-model"})
+
+
 # ──────────────────────────────────────────────────────────────────────
 # Happy paths (mock AutoTokenizer + _detect_vllm_architecture)
 # ──────────────────────────────────────────────────────────────────────
@@ -220,10 +309,11 @@ def _run_prepare(
     fake_tok,
     arch="NeMoSpeechLMForConditionalGeneration",
     llm_arch="Qwen2ForCausalLM",
+    backbone_vocab_size=100,
 ):
     output_dir = _seed_output_dir(tmp_path, llm_arch=llm_arch)
     with (
-        patch.object(to_hf, "_detect_vllm_architecture", return_value=arch),
+        patch.object(to_hf, "_detect_vllm_architecture", return_value=(arch, backbone_vocab_size)),
         patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok),
     ):
         to_hf.prepare_for_vllm(
@@ -261,12 +351,43 @@ def test_prepare_for_vllm_skips_add_if_audio_token_already_in_vocab(tmp_path):
     assert fake_tok.add_special_tokens_calls == []
 
 
+def test_prepare_for_vllm_rejects_invalid_audio_token_id(tmp_path):
+    fake_tok = _FakeTokenizer(vocab_tokens=[AUDIO_TOKEN])
+    fake_tok._vocab[AUDIO_TOKEN] = -1
+
+    with pytest.raises(ValueError, match="valid ID"):
+        _run_prepare(tmp_path, fake_tok)
+
+
+def test_prepare_for_vllm_rejects_audio_token_outside_padded_embeddings(tmp_path):
+    tokens = [f"<extra_{i}>" for i in range(11)] + [AUDIO_TOKEN]
+    fake_tok = _FakeTokenizer(vocab_tokens=tokens)
+
+    with pytest.raises(ValueError, match="outside the SpeechLM embedding table"):
+        _run_prepare(tmp_path, fake_tok, backbone_vocab_size=1)
+
+
+def test_prepare_for_vllm_accepts_audio_token_inside_non_boundary_embedding_row(tmp_path):
+    """Qwen-style reserved rows may put the audio token below the model-vocab boundary."""
+    fake_tok = _FakeTokenizer(vocab_tokens=["<pad>", AUDIO_TOKEN, "<extra>"])
+
+    output_dir = _run_prepare(tmp_path, fake_tok, backbone_vocab_size=100)
+
+    cfg = json.loads((output_dir / "config.json").read_text())
+    assert "audio_token_index" not in cfg
+    assert "image_token_index" not in cfg
+
+
 def test_prepare_for_vllm_uses_training_tokenizer_path(tmp_path):
     """Conversion must persist the tokenizer that supplied training token IDs."""
     output_dir = _seed_output_dir(tmp_path)
     fake_tok = _FakeTokenizer()
     with (
-        patch.object(to_hf, "_detect_vllm_architecture", return_value="NeMoSpeechLMForConditionalGeneration"),
+        patch.object(
+            to_hf,
+            "_detect_vllm_architecture",
+            return_value=("NeMoSpeechLMForConditionalGeneration", 100),
+        ),
         patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok) as load_tokenizer,
     ):
         to_hf.prepare_for_vllm(
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 8bbc95aa3d3c..ffc98bb29704 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -20,7 +20,6 @@
 """
 
 import importlib.util
-import logging
 from types import SimpleNamespace
 
 import pytest
@@ -769,31 +768,6 @@ def test_mtp_patch_extends_mtp_model_types(self, monkeypatch):
 
         assert "nemo_speechlm_mtp" in get_args(_spec_mod.MTPModelTypes)
 
-    def test_mtp_patch_keeps_target_plugin_usable_without_vllm_mtp_guard(self, monkeypatch, caplog):
-        """Older vLLM releases should degrade cleanly instead of breaking registration."""
-        import vllm.config.speculative as _spec_mod
-        from vllm.model_executor.models.registry import ModelRegistry
-
-        from nemo.collections.speechlm2.vllm.salm import register
-
-        registered_archs = []
-        original_register_model = ModelRegistry.register_model
-
-        def _record_registration(architecture, model):
-            registered_archs.append(architecture)
-            return original_register_model(architecture, model)
-
-        monkeypatch.setattr(ModelRegistry, "register_model", _record_registration)
-        monkeypatch.delattr(_spec_mod, "MTPModelTypes")
-
-        with caplog.at_level(logging.WARNING, logger="nemo.collections.speechlm2.vllm.salm"):
-            register()
-
-        assert "NeMoSpeechLMForConditionalGeneration" in registered_archs
-        assert "NeMoSpeechLMMTPModel" not in registered_archs
-        assert "NeMoSpeechLMForConditionalGeneration" in ModelRegistry.get_supported_archs()
-        assert "without MTP speculative-decoding support" in caplog.text
-
     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
@@ -1050,83 +1024,29 @@ def test_embed_input_ids_ignores_embeddings_without_placeholder_positions(self):
 
         assert torch.equal(result, torch.zeros(2, 2))
 
-    def test_embed_input_ids_fallback_supports_nested_audio_chunks(self, monkeypatch):
-        """Older vLLM releases without the private merge helper retain nested-input behavior."""
-        import torch
-
-        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
-
-        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
-        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
-        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(4, 2))
-        result = model.embed_input_ids(
-            torch.tensor([0, 1, 2, 3]),
-            multimodal_embeddings=[[torch.ones(1, 2) * 3], [torch.ones(1, 2) * 7]],
-            is_multimodal=torch.tensor([True, False, True, False]),
-        )
-
-        assert torch.equal(result[0], torch.ones(2) * 3)
-        assert torch.equal(result[1], torch.zeros(2))
-        assert torch.equal(result[2], torch.ones(2) * 7)
-        assert torch.equal(result[3], torch.zeros(2))
-
-    def test_embed_input_ids_fallback_preserves_all_text_mask_noop(self, monkeypatch):
-        """The compatibility fallback should match vLLM for an all-False mask."""
-        import torch
-
-        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
-
-        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
-        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
-        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
-        result = model.embed_input_ids(
-            torch.tensor([0, 1]),
-            multimodal_embeddings=[torch.ones(1, 2)],
-            is_multimodal=torch.tensor([False, False]),
-        )
-
-        assert torch.equal(result, torch.zeros(2, 2))
-
-    def test_embed_input_ids_fallback_rejects_embedding_count_mismatch(self, monkeypatch):
-        """The compatibility fallback should report mismatched audio and placeholder counts."""
-        import torch
-
-        from nemo.collections.speechlm2.vllm.salm import mtp as mtp_module
-
-        monkeypatch.setattr(mtp_module, "_vllm_merge_multimodal_embeddings", None)
-        model = object.__new__(mtp_module.NeMoSpeechLMMTP)
-        model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2))
-
-        with pytest.raises(ValueError, match="2 multimodal tokens to 1 placeholders"):
-            model.embed_input_ids(
-                torch.tensor([0, 1]),
-                multimodal_embeddings=[torch.ones(2, 2)],
-                is_multimodal=torch.tensor([True, False]),
-            )
-
-    def test_target_weight_split_excludes_only_salm_mtp_and_extra_state(self):
-        """The target loader should retain LLM weights while routing SALM draft weights away."""
+    def test_target_weight_split_excludes_salm_mtp_and_rejects_bare_mtp(self):
+        """The target loader should route SALM draft weights and reject unsupported native names."""
         import torch
 
         from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
 
         perception_tensor = torch.ones(1)
         llm_tensor = torch.ones(2)
-        bare_mtp_tensor = torch.ones(3)
         perception, llm = NeMoSpeechLMForConditionalGeneration._split_perception_llm(
             [
                 ("perception.encoder.weight", perception_tensor),
                 ("llm.mtp.layers.0.weight", torch.ones(4)),
                 ("llm.model.layers.0.weight", llm_tensor),
-                ("mtp.layers.0.weight", bare_mtp_tensor),
                 ("llm.model.layers.0._extra_state", torch.ones(5)),
             ]
         )
 
         assert perception == {"encoder.weight": perception_tensor}
-        assert [name for name, _ in llm] == ["llm.model.layers.0.weight", "mtp.layers.0.weight"]
+        assert [name for name, _ in llm] == ["llm.model.layers.0.weight"]
         assert llm[0][1] is llm_tensor
-        assert llm[1][1] is bare_mtp_tensor
+
+        with pytest.raises(ValueError, match=r"llm\.mtp\.\*"):
+            NeMoSpeechLMForConditionalGeneration._split_perception_llm([("mtp.layers.0.weight", torch.ones(3))])
 
     def test_mtp_weight_remap_uses_vllm_embedding_alias(self):
         """Exported SpeechLM embeddings must pass NemotronHMTP's name filter."""
@@ -1290,6 +1210,15 @@ def test_image_token_index_is_unserialized_backbone_vocab_boundary(self):
         assert "audio_token_index" not in cfg.to_dict()
         assert "image_token_index" not in cfg.to_dict()
 
+    @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+    def test_legacy_image_token_index_is_validated_after_backbone_load(self):
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=131072)
+        assert cfg.image_token_index == 131072
+        assert "image_token_index" not in cfg.to_dict()
+
+        with pytest.raises(ValueError, match="backbone vocabulary boundary"):
+            NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=42)
+
 
 class _FakeTokenizer:
     def __init__(self):

From 4ad950b574caa25a042e613151df8a8730d483c1 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 17:47:25 -0700
Subject: [PATCH 08/20] fix(speechlm2): make MTP override spawn-pickleable

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 .../speechlm2/vllm/salm/__init__.py           | 122 ++++++++++-------
 .../collections/speechlm2/vllm/salm/config.py |  10 +-
 .../collections/speechlm2/test_vllm_plugin.py | 127 ++++++++++++++++++
 3 files changed, 209 insertions(+), 50 deletions(-)

diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index 4ba88d843e4d..1f13268cec4d 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -24,6 +24,74 @@
 """
 
 _PKG = "nemo.collections.speechlm2.vllm.salm"
+_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = None
+
+
+def _nemo_speechlm_mtp_hf_config_override(hf_config):
+    """Apply the SpeechLM MTP rewrite, then defer unrelated configs to vLLM.
+
+    This function must remain at module scope: vLLM retains it on the draft
+    ``ModelConfig``, which can cross a spawned process boundary. The original
+    vLLM callable stays in process-local module state because binding the
+    replaced static method inside a ``partial`` also makes that method
+    unresolvable by standard pickle.
+    """
+    if hf_config.model_type == "nemo_speechlm":
+        mtp_cfg = getattr(hf_config, "mtp", None)
+        if not isinstance(mtp_cfg, dict):
+            mtp_cfg = {}
+        # Match SALMAutomodel's training defaults exactly: merely retaining a
+        # recipe depth does not enable MTP, while an enabled block with no
+        # explicit depth constructs one logical head.
+        mtp_enabled = bool(mtp_cfg.get("enabled", False))
+        n_predict = mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0)
+        if mtp_enabled and 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(
+                {
+                    # Size of the physical MTP block that vLLM reuses. A
+                    # repeated-layer checkpoint ships one shared head even
+                    # when it was trained for multiple next-token positions,
+                    # so arbitrary inference K values must be multiples of 1.
+                    # Consequently vLLM defaults to K=1 when K is omitted;
+                    # callers should set num_speculative_tokens explicitly.
+                    "n_predict": 1,
+                    # Physical MTP prediction steps to instantiate. Repeated-layer checkpoints
+                    # ship one shared step (one mtp.layers.* module per hybrid-pattern character)
+                    # that is reapplied every speculative iteration, exactly as vLLM drives its
+                    # MTP draft. This also shadows the backbone text_config's
+                    # num_nextn_predict_layers (e.g. 4), which would otherwise trip the
+                    # single-step assert in NemotronHMultiTokenPredictor.
+                    "num_nextn_predict_layers": 1,
+                    "architectures": ["NeMoSpeechLMMTPModel"],
+                }
+            )
+            return hf_config
+
+    global _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
+    if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
+        # A spawn child can import this module while unpickling the function
+        # without running the vLLM plugin hook first. In that case the class
+        # still exposes its native override, which is safe to capture lazily.
+        from vllm.config.speculative import SpeculativeConfig
+
+        current_override = SpeculativeConfig.hf_config_override
+        if current_override is _nemo_speechlm_mtp_hf_config_override:
+            raise RuntimeError("NeMo SpeechLM MTP override was installed without preserving vLLM's original hook.")
+        _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
+    return _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE(hf_config)
+
+
+_nemo_speechlm_mtp_hf_config_override._nemo_speechlm_mtp_override = True
 
 
 def _patch_vllm_for_nemo_speechlm_mtp() -> None:
@@ -57,53 +125,13 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     # Route SpeechLM MTP checkpoints through SpeculativeConfig.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", None)
-                if not isinstance(mtp_cfg, dict):
-                    mtp_cfg = {}
-                # Match SALMAutomodel's training defaults exactly: merely
-                # retaining a recipe depth does not enable MTP, while an enabled
-                # block with no explicit depth constructs one logical head.
-                mtp_enabled = bool(mtp_cfg.get("enabled", False))
-                n_predict = mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0)
-                if mtp_enabled and 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(
-                        {
-                            # Size of the physical MTP block that vLLM reuses. A
-                            # repeated-layer checkpoint ships one shared head even
-                            # when it was trained for multiple next-token positions,
-                            # so arbitrary inference K values must be multiples of 1.
-                            # Consequently vLLM defaults to K=1 when K is omitted;
-                            # callers should set num_speculative_tokens explicitly.
-                            "n_predict": 1,
-                            # Physical MTP prediction steps to instantiate. Repeated-layer checkpoints
-                            # ship one shared step (one mtp.layers.* module per hybrid-pattern character)
-                            # that is reapplied every speculative iteration, exactly as vLLM drives its
-                            # MTP draft. This also shadows the backbone text_config's
-                            # num_nextn_predict_layers (e.g. 4), which would otherwise trip the
-                            # single-step 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)
+        global _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
+        # Preserve the first native hook for the lifetime of this process.
+        # Replacing it during a later registration could capture a third-party
+        # wrapper that already delegates to us, creating an override cycle.
+        if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
+            _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
+        SpeculativeConfig.hf_config_override = staticmethod(_nemo_speechlm_mtp_hf_config_override)
 
     # Register the SpeechLM MTP draft architecture with vLLM.
     from vllm.model_executor.models.registry import ModelRegistry
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index 46bcde9df9e2..68ae4faea519 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -41,8 +41,10 @@
 
 # Historical serving-time headroom above the backbone vocabulary. vLLM builds
 # the target and draft embedding tables at this padded size, and the weight
-# loader zero-pads the smaller training tensors to match. The tokenizer's
-# actual audio-token ID is validated against this bound during export.
+# loader zero-pads the smaller training tensors to match. ``prepare_for_vllm``
+# validates the tokenizer's audio-token ID against this bound during export;
+# the default export flow treats a validation failure as non-fatal and leaves
+# an HF-only checkpoint.
 _SPEECHLM_EMBED_EXTRA_ROWS = 10
 
 
@@ -229,7 +231,9 @@ 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.
+        vLLM supports any sequence of ``"*"`` (attention) and ``"E"`` (MoE),
+        with one physical MTP layer module instantiated per character. Other
+        characters are rejected by vLLM during model construction.
         """
         mtp_cfg = self.__dict__.get("mtp") or {}
         return mtp_cfg.get("hybrid_override_pattern", "*") if isinstance(mtp_cfg, dict) else "*"
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index ffc98bb29704..5586439cd730 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -43,6 +43,11 @@
 }
 
 
+def _identity_hf_config_override(hf_config):
+    """Pickleable target override for vLLM's composed-override API."""
+    return hf_config
+
+
 @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
 class TestNeMoSpeechLMConfig:
     """Tests for NeMoSpeechLMConfig."""
@@ -710,6 +715,15 @@ def test_register_does_not_load_backbone_config(self, monkeypatch):
 class TestMTPPlugin:
     """Tests for NeMo SpeechLM MTP speculative-decoding support."""
 
+    @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
+
+        original_override = salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
+        yield
+        salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = original_override
+
     @pytest.fixture(autouse=True)
     def mock_backbone_config(self, monkeypatch):
         """Keep registration tests independent of Hugging Face network access."""
@@ -789,6 +803,76 @@ def test_patched_override_routes_nemo_mtp_config(self, monkeypatch):
         assert result.n_predict == 1
         assert result.num_nextn_predict_layers == 1
 
+    def test_patched_override_is_pickleable_for_spawned_engine(self, monkeypatch):
+        """The callable retained on draft ModelConfig must survive multiprocessing spawn."""
+        import pickle
+
+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
+        register()
+
+        restored = pickle.loads(pickle.dumps(SpeculativeConfig.hf_config_override))
+        assert restored is salm_module._nemo_speechlm_mtp_hf_config_override
+
+        overrides = [restored]
+        if hasattr(SpeculativeConfig, "compose_draft_hf_overrides"):
+            composed = SpeculativeConfig.compose_draft_hf_overrides(_identity_hf_config_override)
+            assert composed is not SpeculativeConfig.hf_config_override
+            overrides.append(pickle.loads(pickle.dumps(composed)))
+
+        for override in overrides:
+            hf_cfg = self._HFConfigLike(
+                model_type="nemo_speechlm",
+                mtp={"enabled": True, "num_nextn_predict_layers": 4, "use_repeated_layer": True},
+            )
+            result = override(hf_cfg)
+            assert result.model_type == "nemo_speechlm_mtp"
+            assert result.n_predict == 1
+
+    def test_patched_override_lazily_captures_native_hook_in_spawn_child(self, monkeypatch):
+        """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
+
+        original_calls = []
+
+        def _recording_native(cfg):
+            original_calls.append(cfg)
+            return cfg
+
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
+        monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_native))
+
+        hf_cfg = self._HFConfigLike(model_type="unrelated")
+        result = salm_module._nemo_speechlm_mtp_hf_config_override(hf_cfg)
+
+        assert result is hf_cfg
+        assert original_calls == [hf_cfg]
+        assert salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is _recording_native
+
+    def test_patched_override_rejects_missing_native_hook(self, monkeypatch):
+        """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
+
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
+        monkeypatch.setattr(
+            SpeculativeConfig,
+            "hf_config_override",
+            staticmethod(salm_module._nemo_speechlm_mtp_hf_config_override),
+        )
+
+        with pytest.raises(RuntimeError, match="without preserving vLLM's original hook"):
+            salm_module._nemo_speechlm_mtp_hf_config_override(self._HFConfigLike(model_type="unrelated"))
+
     def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
         """An enabled training block without an explicit depth constructs one head."""
         from transformers import AutoConfig
@@ -830,6 +914,8 @@ def test_patched_override_no_mtp_falls_through(self, monkeypatch):
         from transformers import AutoConfig
         from vllm.config.speculative import SpeculativeConfig
 
+        import nemo.collections.speechlm2.vllm.salm as salm_module
+
         from nemo.collections.speechlm2.vllm.salm import register
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -839,6 +925,7 @@ def _recording_orig(cfg):
             original_calls.append(cfg)
             return cfg
 
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
         monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig))
         register()
 
@@ -852,6 +939,8 @@ def test_patched_override_depth_without_enabled_flag_falls_through(self, monkeyp
         from transformers import AutoConfig
         from vllm.config.speculative import SpeculativeConfig
 
+        import nemo.collections.speechlm2.vllm.salm as salm_module
+
         from nemo.collections.speechlm2.vllm.salm import register
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -861,6 +950,7 @@ def _recording_orig(cfg):
             original_calls.append(cfg)
             return cfg
 
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
         monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig))
         register()
 
@@ -878,6 +968,8 @@ def test_patched_override_explicitly_disabled_mtp_falls_through(self, monkeypatc
         from transformers import AutoConfig
         from vllm.config.speculative import SpeculativeConfig
 
+        import nemo.collections.speechlm2.vllm.salm as salm_module
+
         from nemo.collections.speechlm2.vllm.salm import register
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -887,6 +979,7 @@ def _recording_orig(cfg):
             original_calls.append(cfg)
             return cfg
 
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
         monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig))
         register()
 
@@ -930,6 +1023,40 @@ def test_mtp_override_registration_is_idempotent(self, monkeypatch):
 
         assert SpeculativeConfig.hf_config_override is first_override
 
+    def test_mtp_override_reregistration_preserves_first_native_hook(self, monkeypatch):
+        """A later wrapper that delegates to us must not become our fallback and recurse."""
+        from transformers import AutoConfig
+        from vllm.config.speculative import SpeculativeConfig
+
+        import nemo.collections.speechlm2.vllm.salm as salm_module
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
+        native_calls = []
+
+        def _recording_native(cfg):
+            native_calls.append(cfg)
+            return cfg
+
+        monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
+        monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_native))
+        register()
+        first_override = SpeculativeConfig.hf_config_override
+
+        def _third_party_wrapper(cfg):
+            return first_override(cfg)
+
+        monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_third_party_wrapper))
+        register()
+
+        hf_cfg = self._HFConfigLike(model_type="unrelated")
+        result = SpeculativeConfig.hf_config_override(hf_cfg)
+
+        assert result is hf_cfg
+        assert native_calls == [hf_cfg]
+        assert salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is _recording_native
+
     def test_embed_input_ids_text_only(self):
         """embed_input_ids with no audio embeddings should return plain text embeddings."""
         import torch

From 5bbfd2f353006078cfce066517bf41f8c057c060 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 00:29:28 -0700
Subject: [PATCH 09/20] fix(speechlm2): align export trust and test imports

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   |  3 ++
 tests/collections/speechlm2/test_to_hf.py     |  9 +++++
 .../collections/speechlm2/test_vllm_plugin.py | 36 ++-----------------
 3 files changed, 15 insertions(+), 33 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index a4ecdb156422..289e45937b8c 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -102,6 +102,9 @@ def _canonical_torch_dtype_name(dtype: str | torch.dtype) -> str:
 def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[str, Any]:
     """Build the exported root config without mutating the training config."""
     config = OmegaConf.to_container(model.cfg) if isinstance(model.cfg, DictConfig) else deepcopy(model.cfg)
+    # Remote-code trust is a runtime security decision. Do not persist a
+    # training-time opt-in in checkpoints that may be loaded by another user.
+    config.pop("trust_remote_code", None)
     mtp_cfg = config.get("mtp")
     llm_config = getattr(getattr(model, "llm", None), "config", None)
     actual_mtp_pattern = getattr(llm_config, "mtp_hybrid_override_pattern", None)
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index 45ad325abb9d..edcd7a33093d 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -146,6 +146,15 @@ def test_hf_export_config_persists_built_mtp_pattern_without_mutating_recipe():
     assert model.cfg["mtp"]["hybrid_override_pattern"] == "*"
 
 
+def test_hf_export_config_does_not_persist_remote_code_trust():
+    model = SimpleNamespace(cfg={"trust_remote_code": True})
+
+    config = to_hf._hf_export_config(model, "bfloat16")
+
+    assert "trust_remote_code" not in config
+    assert model.cfg["trust_remote_code"] is True
+
+
 def test_hf_export_config_persists_built_non_repeated_mtp_depth():
     """A preserved native head's actual depth must override stale recipe metadata."""
     model = SimpleNamespace(
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 5586439cd730..81633a49bd18 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -25,9 +25,11 @@
 import pytest
 
 try:
-    from nemo.collections.speechlm2.vllm.salm import config as _config_module
+    import nemo.collections.speechlm2.vllm.salm as _salm_module
+    import nemo.collections.speechlm2.vllm.salm.config as _config_module
 
     NeMoSpeechLMConfig = _config_module.NeMoSpeechLMConfig
+    register = _salm_module.register
 
     _HAS_CONFIG = True
 except (ImportError, RuntimeError):
@@ -648,8 +650,6 @@ def test_register_config(self, monkeypatch):
         """register() should add nemo_speechlm to vLLM's config registry."""
         from transformers import AutoConfig
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(
             AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
         )
@@ -668,8 +668,6 @@ def test_register_model(self, monkeypatch):
         """
         from transformers import AutoConfig
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(
             AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
         )
@@ -686,8 +684,6 @@ def test_register_model(self, monkeypatch):
     def test_register_does_not_patch_fast_tokenizer(self, monkeypatch):
         from transformers import AutoConfig, PreTrainedTokenizerFast
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(
             AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
         )
@@ -701,8 +697,6 @@ def test_register_does_not_load_backbone_config(self, monkeypatch):
 
         from transformers import AutoConfig
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         from_pretrained = Mock(side_effect=AssertionError("register() must not load remote backbone configs"))
         monkeypatch.setattr(AutoConfig, "from_pretrained", from_pretrained)
 
@@ -759,8 +753,6 @@ def test_mtp_patch_registers_model(self, monkeypatch):
         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()
@@ -774,8 +766,6 @@ def test_mtp_patch_extends_mtp_model_types(self, monkeypatch):
         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()
@@ -787,8 +777,6 @@ def test_patched_override_routes_nemo_mtp_config(self, monkeypatch):
         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()
 
@@ -812,8 +800,6 @@ def test_patched_override_is_pickleable_for_spawned_engine(self, monkeypatch):
 
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -878,8 +864,6 @@ def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
         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()
 
@@ -895,8 +879,6 @@ def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeyp
         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()
 
@@ -916,8 +898,6 @@ def test_patched_override_no_mtp_falls_through(self, monkeypatch):
 
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         original_calls = []
 
@@ -941,8 +921,6 @@ def test_patched_override_depth_without_enabled_flag_falls_through(self, monkeyp
 
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         original_calls = []
 
@@ -970,8 +948,6 @@ def test_patched_override_explicitly_disabled_mtp_falls_through(self, monkeypatc
 
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         original_calls = []
 
@@ -997,8 +973,6 @@ def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeyp
         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()
 
@@ -1014,8 +988,6 @@ def test_mtp_override_registration_is_idempotent(self, monkeypatch):
         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
@@ -1030,8 +1002,6 @@ def test_mtp_override_reregistration_preserves_first_native_hook(self, monkeypat
 
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
-        from nemo.collections.speechlm2.vllm.salm import register
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         native_calls = []
 

From 33957a2efdfae073d4bae5d8c06ebaeea1816001 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 00:43:34 -0700
Subject: [PATCH 10/20] fix(tests): normalize SpeechLM config import

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 tests/collections/speechlm2/test_vllm_plugin.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 81633a49bd18..7ace12c67b65 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -117,9 +117,7 @@ def test_hybrid_backbone_aliases_for_vllm(self):
     )
     def test_is_hybrid_backend_helper(self, architectures, expected_is_hybrid):
         """``_is_hybrid_backend`` should match the documented hybrid allow-list."""
-        from nemo.collections.speechlm2.vllm.salm.config import _is_hybrid_backend
-
-        assert _is_hybrid_backend(architectures) is expected_is_hybrid
+        assert _config_module._is_hybrid_backend(architectures) is expected_is_hybrid
 
     @pytest.mark.parametrize(
         "backbone_archs, expected_is_hybrid",

From ecf30a936ddaeb511c205b761b5a4c79fbc44bb3 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 00:46:30 -0700
Subject: [PATCH 11/20] fix(speechlm2): normalize speculative config imports

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 .../speechlm2/vllm/salm/__init__.py           |  9 +++---
 .../collections/speechlm2/test_vllm_plugin.py | 30 ++++---------------
 2 files changed, 9 insertions(+), 30 deletions(-)

diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index 1f13268cec4d..40943c3d8e1b 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -82,9 +82,9 @@ def _nemo_speechlm_mtp_hf_config_override(hf_config):
         # A spawn child can import this module while unpickling the function
         # without running the vLLM plugin hook first. In that case the class
         # still exposes its native override, which is safe to capture lazily.
-        from vllm.config.speculative import SpeculativeConfig
+        import vllm.config.speculative as _spec_mod
 
-        current_override = SpeculativeConfig.hf_config_override
+        current_override = _spec_mod.SpeculativeConfig.hf_config_override
         if current_override is _nemo_speechlm_mtp_hf_config_override:
             raise RuntimeError("NeMo SpeechLM MTP override was installed without preserving vLLM's original hook.")
         _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
@@ -115,7 +115,6 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
     from typing import Literal, get_args
 
     import vllm.config.speculative as _spec_mod
-    from vllm.config.speculative import SpeculativeConfig
 
     # Extend vLLM's recognized MTP model types.
     old_args = get_args(_spec_mod.MTPModelTypes)
@@ -123,7 +122,7 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
         _spec_mod.MTPModelTypes = Literal[old_args + ("nemo_speechlm_mtp",)]
 
     # Route SpeechLM MTP checkpoints through SpeculativeConfig.hf_config_override.
-    current_override = SpeculativeConfig.hf_config_override
+    current_override = _spec_mod.SpeculativeConfig.hf_config_override
     if not getattr(current_override, "_nemo_speechlm_mtp_override", False):
         global _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
         # Preserve the first native hook for the lifetime of this process.
@@ -131,7 +130,7 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
         # wrapper that already delegates to us, creating an override cycle.
         if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
             _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
-        SpeculativeConfig.hf_config_override = staticmethod(_nemo_speechlm_mtp_hf_config_override)
+        _spec_mod.SpeculativeConfig.hf_config_override = staticmethod(_nemo_speechlm_mtp_hf_config_override)
 
     # Register the SpeechLM MTP draft architecture with vLLM.
     from vllm.model_executor.models.registry import ModelRegistry
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 7ace12c67b65..67f91629e5c0 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -36,6 +36,11 @@
     _HAS_CONFIG = False
 
 _HAS_VLLM = importlib.util.find_spec("vllm") is not None
+if _HAS_VLLM:
+    import vllm.config.speculative as _spec_mod
+
+    SpeculativeConfig = _spec_mod.SpeculativeConfig
+
 _DEFAULT_CONFIG_KWARGS = {
     "pretrained_llm": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
     "pretrained_asr": "nvidia/canary-1b-v2",
@@ -761,7 +766,6 @@ 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
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -773,8 +777,6 @@ def test_mtp_patch_extends_mtp_model_types(self, monkeypatch):
     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
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -794,8 +796,6 @@ def test_patched_override_is_pickleable_for_spawned_engine(self, monkeypatch):
         import pickle
 
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -821,8 +821,6 @@ def test_patched_override_is_pickleable_for_spawned_engine(self, monkeypatch):
 
     def test_patched_override_lazily_captures_native_hook_in_spawn_child(self, monkeypatch):
         """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
 
         original_calls = []
@@ -843,8 +841,6 @@ def _recording_native(cfg):
 
     def test_patched_override_rejects_missing_native_hook(self, monkeypatch):
         """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
 
         monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None)
@@ -860,8 +856,6 @@ def test_patched_override_rejects_missing_native_hook(self, monkeypatch):
     def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
         """An enabled training block without an explicit depth constructs one head."""
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -875,8 +869,6 @@ def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
     def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeypatch):
         """Repeated-layer training depth must not constrain inference-time K."""
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -892,8 +884,6 @@ def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeyp
     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
-
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -915,8 +905,6 @@ def _recording_orig(cfg):
     def test_patched_override_depth_without_enabled_flag_falls_through(self, monkeypatch):
         """A retained recipe depth must not enable a head that training did not construct."""
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -942,8 +930,6 @@ def _recording_orig(cfg):
     def test_patched_override_explicitly_disabled_mtp_falls_through(self, monkeypatch):
         """An exported recipe depth must not override mtp.enabled=false."""
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
@@ -969,8 +955,6 @@ def _recording_orig(cfg):
     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
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -984,8 +968,6 @@ def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeyp
     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
-
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
         first_override = SpeculativeConfig.hf_config_override
@@ -996,8 +978,6 @@ def test_mtp_override_registration_is_idempotent(self, monkeypatch):
     def test_mtp_override_reregistration_preserves_first_native_hook(self, monkeypatch):
         """A later wrapper that delegates to us must not become our fallback and recurse."""
         from transformers import AutoConfig
-        from vllm.config.speculative import SpeculativeConfig
-
         import nemo.collections.speechlm2.vllm.salm as salm_module
 
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))

From c5ff50239403d9719bd432295701f46a9828610c Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 00:54:20 -0700
Subject: [PATCH 12/20] style(tests): satisfy black formatting

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 tests/collections/speechlm2/test_vllm_plugin.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 67f91629e5c0..fce7afb9a6f9 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -777,6 +777,7 @@ def test_mtp_patch_extends_mtp_model_types(self, monkeypatch):
     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
+
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -856,6 +857,7 @@ def test_patched_override_rejects_missing_native_hook(self, monkeypatch):
     def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
         """An enabled training block without an explicit depth constructs one head."""
         from transformers import AutoConfig
+
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -869,6 +871,7 @@ def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch):
     def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeypatch):
         """Repeated-layer training depth must not constrain inference-time K."""
         from transformers import AutoConfig
+
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -955,6 +958,7 @@ def _recording_orig(cfg):
     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
+
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
 
@@ -968,6 +972,7 @@ def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeyp
     def test_mtp_override_registration_is_idempotent(self, monkeypatch):
         """register() should not repeatedly wrap the config override."""
         from transformers import AutoConfig
+
         monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()))
         register()
         first_override = SpeculativeConfig.hf_config_override

From 67639bad966631a0ff1bb66014acb87741cb4090 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Wed, 26 Aug 2026 14:39:21 -0700
Subject: [PATCH 13/20] fix(speechlm2): use bundled vLLM backbone config

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   | 13 ++++++-
 .../collections/speechlm2/vllm/salm/config.py | 14 ++++++-
 tests/collections/speechlm2/test_to_hf.py     | 37 +++++++++++++++++++
 .../collections/speechlm2/test_vllm_plugin.py | 32 ++++++++++++++++
 4 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index 289e45937b8c..a5d6159ee8dc 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -232,14 +232,25 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     # 1. Patch config.json (arch, model_type, audio_locator_tag for vLLM plugin).
     arch_model_cfg = dict(model_cfg)
     llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
-    if (llm_backbone_dir / "config.json").exists():
+    llm_backbone_config_path = llm_backbone_dir / "config.json"
+    llm_backbone_config = None
+    if llm_backbone_config_path.exists():
         arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
+        llm_backbone_config = json.loads(llm_backbone_config_path.read_text())
     arch, base_vocab_size = _detect_vllm_architecture(arch_model_cfg)
     config_path = output_dir / "config.json"
     config = json.loads(config_path.read_text())
     config["model_type"] = "nemo_speechlm"
     config["architectures"] = [arch]
     config["audio_locator_tag"] = audio_token
+    if llm_backbone_config is not None:
+        # Keep the export portable while making the bundled config authoritative.
+        # NeMo's HF loader resolves this relative marker to a cached local path;
+        # the vLLM config consumes the embedded copy without another Hub lookup.
+        config["pretrained_llm"] = LLM_BACKBONE_DIR
+        config["llm_config"] = llm_backbone_config
+    else:
+        config.pop("llm_config", None)
     config.pop("audio_token_index", None)
     config.pop("image_token_index", None)
 
diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py
index 68ae4faea519..1ba78a7bd2dd 100644
--- a/nemo/collections/speechlm2/vllm/salm/config.py
+++ b/nemo/collections/speechlm2/vllm/salm/config.py
@@ -74,6 +74,7 @@ def __init__(
         self,
         perception: dict | None = None,
         pretrained_llm: str | None = None,
+        llm_config: dict | None = None,
         pretrained_asr: str | None = None,
         audio_locator_tag: str | None = None,
         prompt_format: str | None = None,
@@ -113,6 +114,7 @@ def __init__(
             # path inert; real checkpoint loads continue through validation below.
             self.perception = {}
             self.pretrained_llm = None
+            self.llm_config = None
             self.pretrained_asr = None
             self.audio_locator_tag = None
             self.prompt_format = None
@@ -140,6 +142,7 @@ def __init__(
             )
         self.perception = perception or {}
         self.pretrained_llm = pretrained_llm
+        self.llm_config = llm_config
         self.pretrained_asr = pretrained_asr
         self.audio_locator_tag = audio_locator_tag
         self.prompt_format = prompt_format
@@ -147,7 +150,16 @@ def __init__(
         self.lora = lora
         self.encoder_chunk_size_seconds = encoder_chunk_size_seconds
 
-        self.text_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True)
+        if llm_config is None:
+            self.text_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True)
+        else:
+            if not isinstance(llm_config, dict):
+                raise ValueError(f"NeMo SpeechLM llm_config must be a dict, got {type(llm_config).__name__}.")
+            embedded_config = dict(llm_config)
+            model_type = embedded_config.pop("model_type", None)
+            if not model_type:
+                raise ValueError("NeMo SpeechLM llm_config must declare model_type.")
+            self.text_config = AutoConfig.for_model(model_type, **embedded_config)
 
         raw_archs = getattr(self.text_config, "architectures", [])
         if len(raw_archs) != 1:
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index edcd7a33093d..5d9f0fbf5eaa 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -345,6 +345,43 @@ def test_prepare_for_vllm_patches_config_json(tmp_path):
     assert cfg["hidden_size"] == 2048
 
 
+def test_prepare_for_vllm_embeds_bundled_backbone_config(tmp_path):
+    """The vLLM root config must not reload the stale training backbone reference."""
+    output_dir = _seed_output_dir(tmp_path)
+    backbone_dir = output_dir / "llm_backbone"
+    backbone_dir.mkdir()
+    backbone_config = {
+        "model_type": "qwen2",
+        "architectures": ["Qwen2ForCausalLM"],
+        "hidden_size": 2048,
+        "vocab_size": 100,
+    }
+    (backbone_dir / "config.json").write_text(json.dumps(backbone_config))
+
+    with (
+        patch.object(
+            to_hf,
+            "_detect_vllm_architecture",
+            return_value=("NeMoSpeechLMForConditionalGeneration", 100),
+        ) as detect_architecture,
+        patch("transformers.AutoTokenizer.from_pretrained", return_value=_FakeTokenizer()),
+    ):
+        to_hf.prepare_for_vllm(
+            str(output_dir),
+            {"pretrained_llm": "stale-training-model", "audio_locator_tag": AUDIO_TOKEN},
+        )
+
+    cfg = json.loads((output_dir / "config.json").read_text())
+    assert cfg["pretrained_llm"] == "llm_backbone"
+    assert cfg["llm_config"] == backbone_config
+    detect_architecture.assert_called_once_with(
+        {
+            "pretrained_llm": str(backbone_dir),
+            "audio_locator_tag": AUDIO_TOKEN,
+        }
+    )
+
+
 def test_prepare_for_vllm_adds_audio_token_to_vocab(tmp_path):
     """Audio token is registered via add_special_tokens when not already in vocab."""
     fake_tok = _FakeTokenizer(vocab_tokens=["<|im_start|>", "<|im_end|>"])
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index fce7afb9a6f9..7f993b875020 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -103,6 +103,38 @@ def test_loads_text_config(self):
         assert hasattr(cfg.text_config, "hidden_size")
         assert cfg.get_text_config() is cfg.text_config
 
+    def test_uses_embedded_backbone_config(self, monkeypatch):
+        """Bundled exports must not reload the stale training backbone reference."""
+        embedded_config = {
+            "model_type": "qwen2",
+            "architectures": ["Qwen2ForCausalLM"],
+            "hidden_size": 2048,
+            "vocab_size": 151936,
+            "num_hidden_layers": 4,
+            "rms_norm_eps": 1e-6,
+        }
+        text_config = SimpleNamespace(**embedded_config)
+        from_pretrained = pytest.fail
+
+        monkeypatch.setattr(
+            _config_module.AutoConfig,
+            "from_pretrained",
+            lambda *args, **kwargs: from_pretrained("must not load pretrained_llm when llm_config is embedded"),
+        )
+        monkeypatch.setattr(_config_module.AutoConfig, "for_model", lambda model_type, **kwargs: text_config)
+
+        cfg = NeMoSpeechLMConfig(
+            **{
+                **_DEFAULT_CONFIG_KWARGS,
+                "pretrained_llm": "llm_backbone",
+                "llm_config": embedded_config,
+            }
+        )
+
+        assert cfg.pretrained_llm == "llm_backbone"
+        assert cfg.llm_config == embedded_config
+        assert cfg.text_config is text_config
+
     def test_hybrid_backbone_aliases_for_vllm(self):
         cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
         assert cfg.is_hybrid is True

From 2e4a4a919e1543a6babcb10db7589a36fe52f595 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Fri, 21 Aug 2026 16:10:57 -0700
Subject: [PATCH 14/20] Add DFlash support to SpeechLM vLLM

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 docs/source/speechlm2/intro.rst               |  1 +
 docs/source/speechlm2/vllm_dflash.rst         | 41 ++++++++++++++
 nemo/collections/speechlm2/vllm/salm/model.py | 26 ++++++++-
 .../collections/speechlm2/test_vllm_plugin.py | 53 +++++++++++++++++++
 4 files changed, 119 insertions(+), 2 deletions(-)
 create mode 100644 docs/source/speechlm2/vllm_dflash.rst

diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst
index 8aac63f9f12f..9fb805304f97 100644
--- a/docs/source/speechlm2/intro.rst
+++ b/docs/source/speechlm2/intro.rst
@@ -393,3 +393,4 @@ For more information, see additional sections in the SpeechLM2 docs:
    datasets
    configs
    training_and_scaling
+   vllm_dflash
diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
new file mode 100644
index 000000000000..0cebe65c5cdd
--- /dev/null
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -0,0 +1,41 @@
+DFlash speculative decoding with vLLM
+======================================
+
+The NeMo SpeechLM vLLM plugin supports checkpoint-backed DFlash speculative
+decoding. DFlash uses intermediate hidden states from the SpeechLM language
+tower to condition a separate draft model; generated draft tokens are verified
+by the target model, so accepted output remains lossless relative to the target.
+
+Requirements
+------------
+
+* A vLLM-ready NeMo SpeechLM checkpoint whose language backbone is compatible
+  with the DFlash draft.
+* vLLM 0.27.1 or later for the published Nemotron 3.5 Lightning recipe.
+* An attention backend that supports the draft model's non-causal attention.
+
+The following example uses the published NVFP4 DFlash draft for the Nemotron
+3.5 Lightning 30B-A3B backbone and proposes six tokens per decoding step:
+
+.. code-block:: bash
+
+   vllm serve /path/to/vllm-ready-speechlm-checkpoint \
+     --trust-remote-code \
+     --speculative-config '{
+       "method": "dflash",
+       "model": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash",
+       "num_speculative_tokens": 6
+     }'
+
+The target SpeechLM checkpoint must use
+``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`` as its language backbone.
+The draft checkpoint provides the auxiliary target-layer selection and mask
+token configuration consumed by vLLM; no draft weights are bundled with NeMo.
+
+Validation
+----------
+
+Compare greedy generation with and without ``--speculative-config`` using the
+same text and audio prompts. The generated token IDs must match. Also inspect
+vLLM's speculative-decoding metrics to confirm that draft tokens are proposed
+and accepted; matching output alone does not prove that DFlash was active.
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index baa97a2b1dde..86073d7f2ab6 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -22,7 +22,8 @@
 ``backends.py`` and is selected once at ``__init__`` time via
 ``make_backend(config)``. The class declares ``IsHybrid`` /
 ``SupportsMambaPrefixCaching`` so vLLM's hybrid KV-cache allocator picks up
-NemotronH backbones; for transformer backbones the runtime
+NemotronH backbones, and ``SupportsEagle3`` so DFlash can consume auxiliary
+hidden states from the language tower. For transformer backbones the runtime
 ``ModelConfig.is_hybrid`` property returns False because ``config.py``
 populates ``text_config.layer_types`` with all-attention markers (vLLM's
 granite-4.0-micro escape hatch).
@@ -40,6 +41,7 @@
 from vllm.model_executor.models.interfaces import (
     IsHybrid,
     MultiModalEmbeddings,
+    SupportsEagle3,
     SupportsMambaPrefixCaching,
     SupportsMultiModal,
     SupportsPP,
@@ -78,6 +80,7 @@ class NeMoSpeechLMForConditionalGeneration(
     SupportsPP,
     IsHybrid,
     SupportsMambaPrefixCaching,
+    SupportsEagle3,
 ):
     """Backbone-agnostic NeMo SpeechLM. Composition with a backend handles per-backbone details."""
 
@@ -112,6 +115,25 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
+    # ── language-model integration ──
+
+    def get_language_model(self) -> nn.Module:
+        """Return the wrapped decoder used by vLLM speculative decoders.
+
+        DFlash resolves the target embedding table and LM head through this
+        hook. Returning the registered vLLM language tower also lets the
+        ``SupportsEagle3`` interface reach its inner ``EagleModelMixin``.
+        """
+        return self.language_model
+
+    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        """Select target layers whose hidden states are consumed by DFlash."""
+        self.language_model.set_aux_hidden_state_layers(layers)
+
+    def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        """Delegate vLLM's fallback auxiliary-layer selection to the decoder."""
+        return self.language_model.get_eagle3_default_aux_hidden_state_layers()
+
     # ── audio processing ──
 
     def _parse_audio_input(
@@ -189,7 +211,7 @@ def forward(
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         if intermediate_tensors is not None:
             inputs_embeds = None
         return self.language_model(input_ids, positions, intermediate_tensors, inputs_embeds)
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 7f993b875020..a957f2b0b4b4 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -740,6 +740,59 @@ 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 TestDFlashPlugin:
+    """Tests for the target-model contract required by real DFlash."""
+
+    def test_model_advertises_eagle3_support(self):
+        from vllm.model_executor.models.interfaces import supports_eagle3
+
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        assert supports_eagle3(NeMoSpeechLMForConditionalGeneration)
+
+    def test_get_language_model_exposes_wrapped_decoder(self):
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        model = object.__new__(NeMoSpeechLMForConditionalGeneration)
+        language_model = object()
+        object.__setattr__(model, "language_model", language_model)
+
+        assert model.get_language_model() is language_model
+
+    def test_aux_hidden_state_methods_delegate_to_wrapped_decoder(self):
+        from unittest.mock import Mock
+
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        layers = (2, 6, 20, 30, 42, 52)
+        language_model = Mock()
+        language_model.get_eagle3_default_aux_hidden_state_layers.return_value = layers
+
+        model = object.__new__(NeMoSpeechLMForConditionalGeneration)
+        object.__setattr__(model, "language_model", language_model)
+
+        model.set_aux_hidden_state_layers(layers)
+
+        language_model.set_aux_hidden_state_layers.assert_called_once_with(layers)
+        assert model.get_eagle3_default_aux_hidden_state_layers() == layers
+
+    def test_forward_preserves_auxiliary_hidden_state_output(self):
+        from unittest.mock import Mock
+
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        output = (object(), [object(), object()])
+        language_model = Mock(return_value=output)
+        model = object.__new__(NeMoSpeechLMForConditionalGeneration)
+        object.__setattr__(model, "language_model", language_model)
+
+        input_ids = object()
+        positions = object()
+        assert model.forward(input_ids, positions) is output
+        language_model.assert_called_once_with(input_ids, positions, None, None)
+
+
 @pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed")
 class TestMTPPlugin:
     """Tests for NeMo SpeechLM MTP speculative-decoding support."""

From e8096663a086f90b110b6882ddfdbdce6d947a8d Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 16:20:28 -0700
Subject: [PATCH 15/20] Add DFlash2 bootstrap support for SpeechLM vLLM

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 docs/source/speechlm2/vllm_dflash.rst         |  57 ++-
 nemo/collections/speechlm2/vllm/salm/model.py |  14 +-
 .../speechlm2/convert_dflash_to_dflash2.py    | 326 ++++++++++++++++++
 .../speechlm2/test_dflash2_checkpoint.py      | 281 +++++++++++++++
 .../collections/speechlm2/test_vllm_plugin.py |   2 +-
 5 files changed, 668 insertions(+), 12 deletions(-)
 create mode 100644 scripts/speechlm2/convert_dflash_to_dflash2.py
 create mode 100644 tests/collections/speechlm2/test_dflash2_checkpoint.py

diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
index 0cebe65c5cdd..1b6fa7c4761f 100644
--- a/docs/source/speechlm2/vllm_dflash.rst
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -1,10 +1,11 @@
 DFlash speculative decoding with vLLM
 ======================================
 
-The NeMo SpeechLM vLLM plugin supports checkpoint-backed DFlash speculative
-decoding. DFlash uses intermediate hidden states from the SpeechLM language
-tower to condition a separate draft model; generated draft tokens are verified
-by the target model, so accepted output remains lossless relative to the target.
+The NeMo SpeechLM vLLM plugin supports checkpoint-backed DFlash and DFlash2
+speculative decoding. Both use intermediate hidden states from the SpeechLM
+language tower to condition a separate draft model; generated draft tokens are
+verified by the target model, so accepted output remains lossless relative to
+the target.
 
 Requirements
 ------------
@@ -32,6 +33,54 @@ The target SpeechLM checkpoint must use
 The draft checkpoint provides the auxiliary target-layer selection and mask
 token configuration consumed by vLLM; no draft weights are bundled with NeMo.
 
+DFlash2 bootstrap
+-----------------
+
+DFlash2 adds two-tap dynamic convolutions and a candidate-path selector. Until
+a trained Lightning DFlash2 checkpoint is published, the existing trained
+Lightning DFlash checkpoint can be converted into a functional DFlash2
+bootstrap:
+
+.. code-block:: bash
+
+   python scripts/speechlm2/convert_dflash_to_dflash2.py \
+     nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash \
+     /path/to/lightning-dflash2-bootstrap
+
+The converter preserves the trained draft backbone, initializes both
+convolutions as exact identities, and initializes the selector as a no-op. It
+also preserves an optional ``mask_embedding.pt`` and excludes the new BF16
+modules from ModelOpt quantization metadata. The source must store its weights
+in a single safetensors file. Its rank-256, top-k-16 selector defaults keep the
+bootstrap memory-representative rather than minimizing its footprint. Output
+files use container-readable model-artifact
+permissions (``0755`` directory and ``0644`` files). The bootstrap therefore
+validates the DFlash2 runtime integration but does not claim the acceptance
+improvement of a checkpoint whose DFlash2 parameters were trained.
+
+At the time of writing, DFlash2 requires the vLLM implementation from pull
+request 52816. It uses the same ``method`` value as DFlash; vLLM selects the
+DFlash2 runtime from the draft checkpoint architecture:
+
+.. code-block:: bash
+
+   pip install -U "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/52816/head"
+
+   vllm serve /path/to/vllm-ready-speechlm-checkpoint \
+     --trust-remote-code \
+     --speculative-config '{
+       "method": "dflash",
+       "model": "/path/to/lightning-dflash2-bootstrap",
+       "num_speculative_tokens": 6
+     }'
+
+The generated config declares ``DFlash2DraftModel``. That architecture forces
+vLLM's V2 model runner; vLLM raises an error if another requested feature is
+incompatible with that runner. The runtime derives its convolution block size
+from ``num_speculative_tokens`` (seven positions in the example: one anchor plus
+six draft tokens). The SpeechLM target uses the same ``SupportsEagle3``
+hidden-state contract for both DFlash versions.
+
 Validation
 ----------
 
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index 86073d7f2ab6..d73ce1fa8bf9 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -22,9 +22,9 @@
 ``backends.py`` and is selected once at ``__init__`` time via
 ``make_backend(config)``. The class declares ``IsHybrid`` /
 ``SupportsMambaPrefixCaching`` so vLLM's hybrid KV-cache allocator picks up
-NemotronH backbones, and ``SupportsEagle3`` so DFlash can consume auxiliary
-hidden states from the language tower. For transformer backbones the runtime
-``ModelConfig.is_hybrid`` property returns False because ``config.py``
+NemotronH backbones, and ``SupportsEagle3`` so DFlash and DFlash2 can consume
+auxiliary hidden states from the language tower. For transformer backbones the
+runtime ``ModelConfig.is_hybrid`` property returns False because ``config.py``
 populates ``text_config.layer_types`` with all-attention markers (vLLM's
 granite-4.0-micro escape hatch).
 
@@ -120,14 +120,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def get_language_model(self) -> nn.Module:
         """Return the wrapped decoder used by vLLM speculative decoders.
 
-        DFlash resolves the target embedding table and LM head through this
-        hook. Returning the registered vLLM language tower also lets the
-        ``SupportsEagle3`` interface reach its inner ``EagleModelMixin``.
+        DFlash and DFlash2 resolve the target embedding table and LM head
+        through this hook. Returning the registered vLLM language tower also
+        lets the ``SupportsEagle3`` interface reach its inner ``EagleModelMixin``.
         """
         return self.language_model
 
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
-        """Select target layers whose hidden states are consumed by DFlash."""
+        """Select target layers whose hidden states are consumed by DFlash drafters."""
         self.language_model.set_aux_hidden_state_layers(layers)
 
     def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
diff --git a/scripts/speechlm2/convert_dflash_to_dflash2.py b/scripts/speechlm2/convert_dflash_to_dflash2.py
new file mode 100644
index 000000000000..bfd05bac6280
--- /dev/null
+++ b/scripts/speechlm2/convert_dflash_to_dflash2.py
@@ -0,0 +1,326 @@
+# 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.
+
+"""Bootstrap a DFlash2 checkpoint from a compatible trained DFlash checkpoint.
+
+The public DFlash2 release adds two-tap grouped convolutions to every draft
+layer and a low-rank candidate selector.  This converter preserves the trained
+DFlash backbone, initializes both convolutions as exact identities, and
+initializes the selector as a no-op.  The resulting checkpoint exercises the
+DFlash2 runtime without pretending that the new parameters have been trained.
+The source checkpoint must store its weights in one safetensors file.
+
+Example::
+
+    python scripts/speechlm2/convert_dflash_to_dflash2.py \
+      nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash \
+      /path/to/lightning-dflash2-bootstrap
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any
+
+
+DEFAULT_CONV_GROUP_SIZE = 16
+DEFAULT_CONV_KERNEL_SIZE = 2
+DEFAULT_SELECTOR_RANK = 256
+DEFAULT_SELECTOR_TOP_K = 16
+_UNQUANTIZED_MODULE_PATTERNS = ("*attention_conv*", "*mlp_conv*", "*candidate_selector*")
+_SOURCE_ARCHITECTURE = "DFlashDraftModel"
+_TARGET_ARCHITECTURE = "DFlash2DraftModel"
+
+
+def build_dflash2_config(
+    source_config: dict[str, Any],
+    *,
+    conv_group_size: int = DEFAULT_CONV_GROUP_SIZE,
+    conv_kernel_size: int = DEFAULT_CONV_KERNEL_SIZE,
+    selector_rank: int = DEFAULT_SELECTOR_RANK,
+    selector_top_k: int = DEFAULT_SELECTOR_TOP_K,
+) -> dict[str, Any]:
+    """Return a validated DFlash2 config derived from ``source_config``."""
+    config = copy.deepcopy(source_config)
+    architectures = config.get("architectures") or []
+    if _SOURCE_ARCHITECTURE not in architectures:
+        raise ValueError(f"The source checkpoint must declare DFlashDraftModel; got architectures={architectures!r}.")
+
+    hidden_size = _positive_int(config, "hidden_size")
+    _positive_int(config, "num_hidden_layers")
+    _positive_int(config, "vocab_size")
+    for name, value in (
+        ("conv_group_size", conv_group_size),
+        ("conv_kernel_size", conv_kernel_size),
+        ("selector_rank", selector_rank),
+        ("selector_top_k", selector_top_k),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+            raise ValueError(f"{name} must be a positive integer; got {value!r}.")
+    if conv_group_size > hidden_size or hidden_size % conv_group_size:
+        raise ValueError(f"conv_group_size={conv_group_size} must divide hidden_size={hidden_size}.")
+    if selector_top_k > config["vocab_size"]:
+        raise ValueError(f"selector_top_k={selector_top_k} cannot exceed vocab_size={config['vocab_size']}.")
+
+    dflash_config = dict(config.get("dflash_config") or {})
+    target_layer_ids = dflash_config.get("target_layer_ids") or config.get("target_layer_ids")
+    if not isinstance(target_layer_ids, list) or not target_layer_ids:
+        raise ValueError("The source checkpoint must define non-empty dflash_config.target_layer_ids.")
+    if any(not isinstance(layer, int) or isinstance(layer, bool) or layer < 0 for layer in target_layer_ids):
+        raise ValueError(f"target_layer_ids must contain non-negative integers; got {target_layer_ids!r}.")
+    if "causal" in dflash_config and not isinstance(dflash_config["causal"], bool):
+        raise ValueError(f"dflash_config.causal must be a boolean; got {dflash_config['causal']!r}.")
+
+    dflash_config.update(
+        {
+            "target_layer_ids": target_layer_ids,
+            "conv_group_size": conv_group_size,
+            "conv_kernel_size": conv_kernel_size,
+            "selector_rank": selector_rank,
+            "selector_top_k": selector_top_k,
+        }
+    )
+    config["architectures"] = [_TARGET_ARCHITECTURE]
+    config["dflash_config"] = dflash_config
+    if "is_causal" not in config and "causal" in dflash_config:
+        config["is_causal"] = bool(dflash_config["causal"])
+    config.pop("num_target_layers", None)
+    quantization_config = config.get("quantization_config")
+    if isinstance(quantization_config, dict):
+        for key in ("ignore", "exclude_modules"):
+            _extend_patterns(quantization_config, key)
+    return config
+
+
+def dflash2_tensor_shapes(config: dict[str, Any]) -> dict[str, tuple[int, ...]]:
+    """Describe the additional tensors required by vLLM's DFlash2 model."""
+    hidden_size = _positive_int(config, "hidden_size")
+    num_layers = _positive_int(config, "num_hidden_layers")
+    vocab_size = _positive_int(config, "vocab_size")
+    dflash_config = config.get("dflash_config") or {}
+    group_size = int(dflash_config["conv_group_size"])
+    kernel_size = int(dflash_config["conv_kernel_size"])
+    selector_rank = int(dflash_config["selector_rank"])
+    if hidden_size % group_size:
+        raise ValueError(f"conv_group_size={group_size} must divide hidden_size={hidden_size}.")
+
+    num_groups = hidden_size // group_size
+    shapes: dict[str, tuple[int, ...]] = {}
+    for layer in range(num_layers):
+        for name in ("attention_conv", "mlp_conv"):
+            prefix = f"layers.{layer}.{name}"
+            shapes[f"{prefix}.base_kernel"] = (2, kernel_size, hidden_size)
+            # ReplicatedLinear stores [output, input]. The output rows flatten
+            # vLLM's (side, tap, group) coefficient layout in that order.
+            shapes[f"{prefix}.kernel_projection.weight"] = (
+                2 * kernel_size * num_groups,
+                hidden_size,
+            )
+    shapes["candidate_selector.predecessor_codebook"] = (vocab_size, selector_rank)
+    shapes["candidate_selector.successor_codebook"] = (vocab_size, selector_rank)
+    shapes["candidate_selector.hidden_projection.weight"] = (selector_rank, hidden_size)
+    return shapes
+
+
+def initialize_dflash2_tensors(config: dict[str, Any]):
+    """Create identity convolutions and a neutral selector in BF16.
+
+    DFlash2's unquantized runtime modules load these tensors in the model dtype.
+    BF16 deliberately matches the Lightning target and keeps the bootstrap
+    memory-representative of a trained Lightning DFlash2 checkpoint.
+    """
+    try:
+        import torch
+    except ImportError as error:
+        raise RuntimeError("Checkpoint conversion requires PyTorch.") from error
+
+    tensors = {}
+    for name, shape in dflash2_tensor_shapes(config).items():
+        tensor = torch.zeros(shape, dtype=torch.bfloat16)
+        if name.endswith(".base_kernel"):
+            # Both the pre-attention/MLP and post-attention/MLP convolutions
+            # pass the current token through unchanged. All older taps and all
+            # dynamic coefficients remain zero.
+            tensor[:, 0, :].fill_(1)
+        tensors[name] = tensor
+    return tensors
+
+
+def convert_checkpoint(
+    source: str,
+    output: str | Path,
+    *,
+    conv_group_size: int = DEFAULT_CONV_GROUP_SIZE,
+    conv_kernel_size: int = DEFAULT_CONV_KERNEL_SIZE,
+    selector_rank: int = DEFAULT_SELECTOR_RANK,
+    selector_top_k: int = DEFAULT_SELECTOR_TOP_K,
+) -> Path:
+    """Convert ``source`` into a new local DFlash2 checkpoint directory."""
+    source_dir, source_label = _resolve_source(source)
+    config_path = source_dir / "config.json"
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Missing source config: {config_path}")
+    source_config = json.loads(config_path.read_text())
+    config = build_dflash2_config(
+        source_config,
+        conv_group_size=conv_group_size,
+        conv_kernel_size=conv_kernel_size,
+        selector_rank=selector_rank,
+        selector_top_k=selector_top_k,
+    )
+
+    output_dir = Path(output).expanduser().resolve()
+    if output_dir.exists():
+        raise FileExistsError(f"Output already exists: {output_dir}")
+    output_dir.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}-", dir=output_dir.parent) as staging:
+        staging_dir = Path(staging)
+        _write_checkpoint(source_dir, staging_dir, config, source_label)
+        staging_dir.rename(output_dir)
+    return output_dir
+
+
+def _positive_int(config: dict[str, Any], key: str) -> int:
+    value = config.get(key)
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"config.{key} must be a positive integer; got {value!r}.")
+    return value
+
+
+def _extend_patterns(config: dict[str, Any], key: str) -> None:
+    patterns = config.get(key)
+    if patterns is None:
+        patterns = []
+    if not isinstance(patterns, list) or any(not isinstance(pattern, str) for pattern in patterns):
+        raise ValueError(f"quantization_config.{key} must be a list of strings; got {patterns!r}.")
+    config[key] = [*patterns, *(pattern for pattern in _UNQUANTIZED_MODULE_PATTERNS if pattern not in patterns)]
+
+
+def _resolve_source(source: str) -> tuple[Path, str]:
+    source_path = Path(source).expanduser()
+    if source_path.is_dir():
+        return source_path.resolve(), str(source_path.resolve())
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as error:
+        raise RuntimeError(
+            "Resolving a Hugging Face model ID requires huggingface_hub. "
+            "Install NeMo's standard dependencies or pass a local checkpoint directory."
+        ) from error
+
+    snapshot = snapshot_download(
+        source,
+        allow_patterns=["*.json", "*.safetensors", "mask_embedding.pt", "LICENSE*", "README*"],
+    )
+    return Path(snapshot), source
+
+
+def _copy_metadata(source_dir: Path, output_dir: Path) -> None:
+    for path in source_dir.glob("LICENSE*"):
+        shutil.copy2(path, output_dir / path.name)
+    mask_embedding = source_dir / "mask_embedding.pt"
+    if mask_embedding.is_file():
+        shutil.copy2(mask_embedding, output_dir / mask_embedding.name)
+
+    source_quant_config = source_dir / "hf_quant_config.json"
+    if source_quant_config.is_file():
+        quant_config = json.loads(source_quant_config.read_text())
+        quantization = quant_config.get("quantization")
+        if isinstance(quantization, dict):
+            _extend_patterns(quantization, "exclude_modules")
+        (output_dir / source_quant_config.name).write_text(json.dumps(quant_config, indent=2) + "\n")
+
+
+def _write_checkpoint(source_dir: Path, output_dir: Path, config: dict[str, Any], source_label: str) -> None:
+    try:
+        from safetensors import safe_open
+        from safetensors.torch import save_file
+    except ImportError as error:
+        raise RuntimeError("Checkpoint conversion requires safetensors.") from error
+
+    weight_files = sorted(source_dir.glob("*.safetensors"))
+    if len(weight_files) != 1:
+        raise ValueError(
+            "This converter expects one safetensors file in the source checkpoint; "
+            f"found {[path.name for path in weight_files]!r}."
+        )
+
+    source_weights = weight_files[0]
+    with safe_open(source_weights, framework="pt", device="cpu") as handle:
+        metadata = handle.metadata()
+        tensors = {name: handle.get_tensor(name) for name in handle.keys()}
+
+    additions = initialize_dflash2_tensors(config)
+    overlap = sorted(tensors.keys() & additions.keys())
+    if overlap:
+        raise ValueError(f"Source checkpoint already contains DFlash2 tensors: {overlap[:5]!r}.")
+    tensors.update(additions)
+
+    save_file(tensors, output_dir / "model.safetensors", metadata=metadata)
+    (output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    manifest = {
+        "source": source_label,
+        "initialization": {
+            "convolutions": "identity",
+            "candidate_selector": "neutral",
+        },
+        "trained_dflash2_parameters": False,
+    }
+    (output_dir / "dflash2_bootstrap.json").write_text(json.dumps(manifest, indent=2) + "\n")
+    (output_dir / "README.md").write_text(
+        "# Lightning DFlash2 bootstrap\n\n"
+        f"This checkpoint was derived from `{source_label}`. The trained DFlash "
+        "backbone is preserved, while the DFlash2 convolutions are initialized as "
+        "identities and its candidate selector is neutral. It is a functional "
+        "DFlash2 runtime artifact, not a DFlash2-trained performance release.\n"
+    )
+    _copy_metadata(source_dir, output_dir)
+
+    # TemporaryDirectory and safetensors default to 0700/0600. The converted
+    # public model is commonly bind-mounted into a root-squashed inference
+    # container, so normalize it to read-only model-artifact permissions.
+    output_dir.chmod(0o755)
+    for path in output_dir.iterdir():
+        if path.is_file():
+            path.chmod(0o644)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("source", help="Local DFlash checkpoint directory or Hugging Face model ID")
+    parser.add_argument("output", help="New local output directory")
+    parser.add_argument("--conv-group-size", type=int, default=DEFAULT_CONV_GROUP_SIZE)
+    parser.add_argument("--conv-kernel-size", type=int, default=DEFAULT_CONV_KERNEL_SIZE)
+    parser.add_argument("--selector-rank", type=int, default=DEFAULT_SELECTOR_RANK)
+    parser.add_argument("--selector-top-k", type=int, default=DEFAULT_SELECTOR_TOP_K)
+    args = parser.parse_args()
+    output = convert_checkpoint(
+        args.source,
+        args.output,
+        conv_group_size=args.conv_group_size,
+        conv_kernel_size=args.conv_kernel_size,
+        selector_rank=args.selector_rank,
+        selector_top_k=args.selector_top_k,
+    )
+    print(f"Created DFlash2 checkpoint: {output}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/collections/speechlm2/test_dflash2_checkpoint.py b/tests/collections/speechlm2/test_dflash2_checkpoint.py
new file mode 100644
index 000000000000..6de12b5b7a42
--- /dev/null
+++ b/tests/collections/speechlm2/test_dflash2_checkpoint.py
@@ -0,0 +1,281 @@
+# 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.
+
+import json
+
+import pytest
+
+from scripts.speechlm2.convert_dflash_to_dflash2 import (
+    build_dflash2_config,
+    convert_checkpoint,
+    dflash2_tensor_shapes,
+    initialize_dflash2_tensors,
+)
+
+
+@pytest.fixture
+def lightning_dflash_config():
+    return {
+        "architectures": ["DFlashDraftModel"],
+        "hidden_size": 2688,
+        "num_hidden_layers": 6,
+        "num_target_layers": 52,
+        "vocab_size": 131072,
+        "dflash_config": {
+            "causal": False,
+            "mask_token_id": 990,
+            "target_layer_ids": [1, 5, 19, 29, 41, 51],
+        },
+        "quantization_config": {
+            "ignore": ["*embed_tokens*", "*self_attn*"],
+            "exclude_modules": ["*embed_tokens*", "*self_attn*"],
+        },
+    }
+
+
+@pytest.fixture
+def tiny_dflash_checkpoint(tmp_path, lightning_dflash_config):
+    torch = pytest.importorskip("torch")
+    safetensors = pytest.importorskip("safetensors.torch")
+    source = tmp_path / "source"
+    source.mkdir()
+    config = {
+        **lightning_dflash_config,
+        "hidden_size": 8,
+        "num_hidden_layers": 1,
+        "vocab_size": 32,
+        "dflash_config": {
+            **lightning_dflash_config["dflash_config"],
+            "target_layer_ids": [1],
+        },
+    }
+    (source / "config.json").write_text(json.dumps(config))
+    (source / "hf_quant_config.json").write_text(
+        json.dumps({"quantization": {"exclude_modules": ["*embed_tokens*", "*self_attn*"]}})
+    )
+    (source / "mask_embedding.pt").write_bytes(b"test mask embedding")
+    safetensors.save_file({"norm.weight": torch.arange(8, dtype=torch.bfloat16)}, source / "model.safetensors")
+    return source
+
+
+def test_build_config_is_lightning_compatible(lightning_dflash_config):
+    converted = build_dflash2_config(lightning_dflash_config)
+
+    assert converted["architectures"] == ["DFlash2DraftModel"]
+    assert "num_target_layers" not in converted
+    assert converted["is_causal"] is False
+    assert converted["dflash_config"] == {
+        "causal": False,
+        "mask_token_id": 990,
+        "target_layer_ids": [1, 5, 19, 29, 41, 51],
+        "conv_group_size": 16,
+        "conv_kernel_size": 2,
+        "selector_rank": 256,
+        "selector_top_k": 16,
+    }
+    expected_exclusions = [
+        "*embed_tokens*",
+        "*self_attn*",
+        "*attention_conv*",
+        "*mlp_conv*",
+        "*candidate_selector*",
+    ]
+    assert converted["quantization_config"]["ignore"] == expected_exclusions
+    assert converted["quantization_config"]["exclude_modules"] == expected_exclusions
+
+
+def test_build_config_preserves_layer_type_causality(lightning_dflash_config):
+    del lightning_dflash_config["dflash_config"]["causal"]
+    lightning_dflash_config["layer_types"] = ["sliding_attention", "full_attention"]
+
+    converted = build_dflash2_config(lightning_dflash_config)
+
+    assert "is_causal" not in converted
+    assert converted["layer_types"] == ["sliding_attention", "full_attention"]
+
+
+def test_build_config_normalizes_top_level_target_layers(lightning_dflash_config):
+    target_layer_ids = lightning_dflash_config["dflash_config"].pop("target_layer_ids")
+    lightning_dflash_config["target_layer_ids"] = target_layer_ids
+
+    converted = build_dflash2_config(lightning_dflash_config)
+
+    assert converted["dflash_config"]["target_layer_ids"] == target_layer_ids
+
+
+def test_tensor_shapes_match_vllm_dflash2(lightning_dflash_config):
+    shapes = dflash2_tensor_shapes(build_dflash2_config(lightning_dflash_config))
+
+    assert len(shapes) == 27
+    assert shapes["layers.0.attention_conv.base_kernel"] == (2, 2, 2688)
+    assert shapes["layers.5.mlp_conv.kernel_projection.weight"] == (672, 2688)
+    assert shapes["candidate_selector.predecessor_codebook"] == (131072, 256)
+    assert shapes["candidate_selector.successor_codebook"] == (131072, 256)
+    assert shapes["candidate_selector.hidden_projection.weight"] == (256, 2688)
+
+
+def test_tensor_shapes_match_installed_vllm_modules(monkeypatch, lightning_dflash_config):
+    torch = pytest.importorskip("torch")
+    linear = pytest.importorskip("vllm.model_executor.layers.linear")
+    parameter = pytest.importorskip("vllm.model_executor.parameter")
+    dflash2 = pytest.importorskip("vllm.model_executor.models.qwen3_dflash2")
+    vllm_config = pytest.importorskip("vllm.config")
+
+    for module in (linear, parameter):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    config = {
+        **lightning_dflash_config,
+        "hidden_size": 8,
+        "num_hidden_layers": 1,
+        "vocab_size": 32,
+    }
+    config = build_dflash2_config(config, conv_group_size=4, selector_rank=2, selector_top_k=4)
+    with vllm_config.set_current_vllm_config(vllm_config.VllmConfig()):
+        conv = dflash2.DFlashGroupedConv(
+            hidden_size=8,
+            taps=2,
+            group_size=4,
+            block_size=3,
+            params_dtype=torch.bfloat16,
+            prefix="attention_conv",
+        )
+        selector = dflash2.CandidateSelector(
+            hidden_size=8,
+            vocab_size=32,
+            rank=2,
+            top_k=4,
+            params_dtype=torch.bfloat16,
+            prefix="candidate_selector",
+        )
+
+    actual = {
+        **{f"layers.0.attention_conv.{name}": tuple(param.shape) for name, param in conv.named_parameters()},
+        **{f"candidate_selector.{name}": tuple(param.shape) for name, param in selector.named_parameters()},
+    }
+    expected = dflash2_tensor_shapes(config)
+    covered_expected = {
+        name for name in expected if name.startswith(("layers.0.attention_conv.", "candidate_selector."))
+    }
+    assert set(actual) == covered_expected
+    for name, shape in actual.items():
+        assert shape == expected[name]
+
+
+def test_identity_convolutions_and_neutral_selector(lightning_dflash_config):
+    torch = pytest.importorskip("torch")
+    tensors = initialize_dflash2_tensors(build_dflash2_config(lightning_dflash_config))
+
+    base = tensors["layers.0.attention_conv.base_kernel"]
+    torch.testing.assert_close(base[:, 0], torch.ones_like(base[:, 0]))
+    torch.testing.assert_close(base[:, 1], torch.zeros_like(base[:, 1]))
+    assert not tensors["layers.0.attention_conv.kernel_projection.weight"].count_nonzero()
+    assert not tensors["candidate_selector.predecessor_codebook"].count_nonzero()
+    assert not tensors["candidate_selector.successor_codebook"].count_nonzero()
+    assert not tensors["candidate_selector.hidden_projection.weight"].count_nonzero()
+
+
+def test_convert_checkpoint_writes_loadable_artifact(tmp_path, tiny_dflash_checkpoint):
+    torch = pytest.importorskip("torch")
+    safetensors = pytest.importorskip("safetensors.torch")
+    original = safetensors.load_file(tiny_dflash_checkpoint / "model.safetensors")["norm.weight"]
+
+    output = convert_checkpoint(
+        str(tiny_dflash_checkpoint),
+        tmp_path / "output",
+        conv_group_size=4,
+        selector_rank=2,
+        selector_top_k=4,
+    )
+
+    output_config = json.loads((output / "config.json").read_text())
+    output_weights = safetensors.load_file(output / "model.safetensors")
+    manifest = json.loads((output / "dflash2_bootstrap.json").read_text())
+    hf_quant_config = json.loads((output / "hf_quant_config.json").read_text())
+    assert output_config["architectures"] == ["DFlash2DraftModel"]
+    torch.testing.assert_close(output_weights["norm.weight"], original)
+    assert output_weights["layers.0.attention_conv.kernel_projection.weight"].shape == (8, 8)
+    assert manifest["trained_dflash2_parameters"] is False
+    assert (output / "hf_quant_config.json").is_file()
+    assert (output / "mask_embedding.pt").read_bytes() == b"test mask embedding"
+    assert hf_quant_config["quantization"]["exclude_modules"] == [
+        "*embed_tokens*",
+        "*self_attn*",
+        "*attention_conv*",
+        "*mlp_conv*",
+        "*candidate_selector*",
+    ]
+    assert output.stat().st_mode & 0o777 == 0o755
+    assert (output / "model.safetensors").stat().st_mode & 0o777 == 0o644
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"conv_group_size": 17}, "must divide hidden_size"),
+        ({"selector_top_k": 131073}, "cannot exceed vocab_size"),
+        ({"conv_kernel_size": 0}, "must be a positive integer"),
+    ],
+)
+def test_rejects_incompatible_dimensions(lightning_dflash_config, kwargs, match):
+    with pytest.raises(ValueError, match=match):
+        build_dflash2_config(lightning_dflash_config, **kwargs)
+
+
+def test_rejects_non_dflash_source(lightning_dflash_config):
+    lightning_dflash_config["architectures"] = ["DFlash2DraftModel"]
+
+    with pytest.raises(ValueError, match="must declare DFlashDraftModel"):
+        build_dflash2_config(lightning_dflash_config)
+
+
+def test_rejects_existing_output(tmp_path, tiny_dflash_checkpoint):
+    output = tmp_path / "output"
+    output.mkdir()
+
+    with pytest.raises(FileExistsError, match="Output already exists"):
+        convert_checkpoint(str(tiny_dflash_checkpoint), output, conv_group_size=4, selector_rank=2)
+
+
+def test_rejects_multiple_weight_files(tmp_path, tiny_dflash_checkpoint):
+    (tiny_dflash_checkpoint / "second.safetensors").write_bytes(b"not read")
+
+    with pytest.raises(ValueError, match="expects one safetensors file"):
+        convert_checkpoint(
+            str(tiny_dflash_checkpoint),
+            tmp_path / "output",
+            conv_group_size=4,
+            selector_rank=2,
+        )
+
+
+def test_rejects_existing_dflash2_tensors(tmp_path, tiny_dflash_checkpoint):
+    torch = pytest.importorskip("torch")
+    safetensors = pytest.importorskip("safetensors.torch")
+    safetensors.save_file(
+        {
+            "norm.weight": torch.arange(8, dtype=torch.bfloat16),
+            "layers.0.attention_conv.base_kernel": torch.zeros((2, 2, 8), dtype=torch.bfloat16),
+        },
+        tiny_dflash_checkpoint / "model.safetensors",
+    )
+
+    with pytest.raises(ValueError, match="already contains DFlash2 tensors"):
+        convert_checkpoint(
+            str(tiny_dflash_checkpoint),
+            tmp_path / "output",
+            conv_group_size=4,
+            selector_rank=2,
+        )
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index a957f2b0b4b4..8f491239edd1 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -742,7 +742,7 @@ def test_register_does_not_load_backbone_config(self, monkeypatch):
 
 @pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed")
 class TestDFlashPlugin:
-    """Tests for the target-model contract required by real DFlash."""
+    """Tests for the target-model contract required by DFlash and DFlash2."""
 
     def test_model_advertises_eagle3_support(self):
         from vllm.model_executor.models.interfaces import supports_eagle3

From 06bb29bb91877601472bb722fb367967681bb083 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 16:30:24 -0700
Subject: [PATCH 16/20] Keep DFlash2 support inference-only

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 docs/source/speechlm2/vllm_dflash.rst         |  51 ++-
 .../speechlm2/convert_dflash_to_dflash2.py    | 326 ------------------
 .../speechlm2/test_dflash2_checkpoint.py      | 281 ---------------
 3 files changed, 21 insertions(+), 637 deletions(-)
 delete mode 100644 scripts/speechlm2/convert_dflash_to_dflash2.py
 delete mode 100644 tests/collections/speechlm2/test_dflash2_checkpoint.py

diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
index 1b6fa7c4761f..bd1bcb2860e3 100644
--- a/docs/source/speechlm2/vllm_dflash.rst
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -33,34 +33,17 @@ The target SpeechLM checkpoint must use
 The draft checkpoint provides the auxiliary target-layer selection and mask
 token configuration consumed by vLLM; no draft weights are bundled with NeMo.
 
-DFlash2 bootstrap
+DFlash2 inference
 -----------------
 
-DFlash2 adds two-tap dynamic convolutions and a candidate-path selector. Until
-a trained Lightning DFlash2 checkpoint is published, the existing trained
-Lightning DFlash checkpoint can be converted into a functional DFlash2
-bootstrap:
-
-.. code-block:: bash
-
-   python scripts/speechlm2/convert_dflash_to_dflash2.py \
-     nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash \
-     /path/to/lightning-dflash2-bootstrap
-
-The converter preserves the trained draft backbone, initializes both
-convolutions as exact identities, and initializes the selector as a no-op. It
-also preserves an optional ``mask_embedding.pt`` and excludes the new BF16
-modules from ModelOpt quantization metadata. The source must store its weights
-in a single safetensors file. Its rank-256, top-k-16 selector defaults keep the
-bootstrap memory-representative rather than minimizing its footprint. Output
-files use container-readable model-artifact
-permissions (``0755`` directory and ``0644`` files). The bootstrap therefore
-validates the DFlash2 runtime integration but does not claim the acceptance
-improvement of a checkpoint whose DFlash2 parameters were trained.
+DFlash2 adds dynamic convolutions and a candidate-path selector to the draft
+model. Its checkpoint must be trained or fine-tuned separately for the target
+language backbone; NeMo's vLLM inference plugin does not create or convert
+DFlash2 weights.
 
 At the time of writing, DFlash2 requires the vLLM implementation from pull
 request 52816. It uses the same ``method`` value as DFlash; vLLM selects the
-DFlash2 runtime from the draft checkpoint architecture:
+DFlash2 runtime from the trained draft checkpoint's architecture:
 
 .. code-block:: bash
 
@@ -70,16 +53,24 @@ DFlash2 runtime from the draft checkpoint architecture:
      --trust-remote-code \
      --speculative-config '{
        "method": "dflash",
-       "model": "/path/to/lightning-dflash2-bootstrap",
+       "model": "/path/to/trained-lightning-dflash2-checkpoint",
        "num_speculative_tokens": 6
      }'
 
-The generated config declares ``DFlash2DraftModel``. That architecture forces
-vLLM's V2 model runner; vLLM raises an error if another requested feature is
-incompatible with that runner. The runtime derives its convolution block size
-from ``num_speculative_tokens`` (seven positions in the example: one anchor plus
-six draft tokens). The SpeechLM target uses the same ``SupportsEagle3``
-hidden-state contract for both DFlash versions.
+The trained draft config must declare ``DFlash2DraftModel``. Its
+``dflash_config`` must include ``target_layer_ids``, ``conv_group_size``,
+``conv_kernel_size``, ``selector_rank``, and ``selector_top_k``. Set
+``num_speculative_tokens`` to one less than the convolution block size used to
+train the draft: vLLM constructs each runtime block from one anchor plus the
+configured number of draft tokens and does not reject a training/inference
+block-size mismatch.
+
+The DFlash2 architecture forces vLLM's V2 model runner, including for hybrid
+NemotronH targets that would otherwise use V1. vLLM raises an error for features
+it knows are incompatible with V2, but the target and serving configuration
+should still be qualified on that runner. The SpeechLM target uses the same
+``SupportsEagle3`` hidden-state contract for DFlash and DFlash2, so no draft
+weights or training logic are bundled with NeMo.
 
 Validation
 ----------
diff --git a/scripts/speechlm2/convert_dflash_to_dflash2.py b/scripts/speechlm2/convert_dflash_to_dflash2.py
deleted file mode 100644
index bfd05bac6280..000000000000
--- a/scripts/speechlm2/convert_dflash_to_dflash2.py
+++ /dev/null
@@ -1,326 +0,0 @@
-# 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.
-
-"""Bootstrap a DFlash2 checkpoint from a compatible trained DFlash checkpoint.
-
-The public DFlash2 release adds two-tap grouped convolutions to every draft
-layer and a low-rank candidate selector.  This converter preserves the trained
-DFlash backbone, initializes both convolutions as exact identities, and
-initializes the selector as a no-op.  The resulting checkpoint exercises the
-DFlash2 runtime without pretending that the new parameters have been trained.
-The source checkpoint must store its weights in one safetensors file.
-
-Example::
-
-    python scripts/speechlm2/convert_dflash_to_dflash2.py \
-      nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash \
-      /path/to/lightning-dflash2-bootstrap
-"""
-
-from __future__ import annotations
-
-import argparse
-import copy
-import json
-import shutil
-import tempfile
-from pathlib import Path
-from typing import Any
-
-
-DEFAULT_CONV_GROUP_SIZE = 16
-DEFAULT_CONV_KERNEL_SIZE = 2
-DEFAULT_SELECTOR_RANK = 256
-DEFAULT_SELECTOR_TOP_K = 16
-_UNQUANTIZED_MODULE_PATTERNS = ("*attention_conv*", "*mlp_conv*", "*candidate_selector*")
-_SOURCE_ARCHITECTURE = "DFlashDraftModel"
-_TARGET_ARCHITECTURE = "DFlash2DraftModel"
-
-
-def build_dflash2_config(
-    source_config: dict[str, Any],
-    *,
-    conv_group_size: int = DEFAULT_CONV_GROUP_SIZE,
-    conv_kernel_size: int = DEFAULT_CONV_KERNEL_SIZE,
-    selector_rank: int = DEFAULT_SELECTOR_RANK,
-    selector_top_k: int = DEFAULT_SELECTOR_TOP_K,
-) -> dict[str, Any]:
-    """Return a validated DFlash2 config derived from ``source_config``."""
-    config = copy.deepcopy(source_config)
-    architectures = config.get("architectures") or []
-    if _SOURCE_ARCHITECTURE not in architectures:
-        raise ValueError(f"The source checkpoint must declare DFlashDraftModel; got architectures={architectures!r}.")
-
-    hidden_size = _positive_int(config, "hidden_size")
-    _positive_int(config, "num_hidden_layers")
-    _positive_int(config, "vocab_size")
-    for name, value in (
-        ("conv_group_size", conv_group_size),
-        ("conv_kernel_size", conv_kernel_size),
-        ("selector_rank", selector_rank),
-        ("selector_top_k", selector_top_k),
-    ):
-        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
-            raise ValueError(f"{name} must be a positive integer; got {value!r}.")
-    if conv_group_size > hidden_size or hidden_size % conv_group_size:
-        raise ValueError(f"conv_group_size={conv_group_size} must divide hidden_size={hidden_size}.")
-    if selector_top_k > config["vocab_size"]:
-        raise ValueError(f"selector_top_k={selector_top_k} cannot exceed vocab_size={config['vocab_size']}.")
-
-    dflash_config = dict(config.get("dflash_config") or {})
-    target_layer_ids = dflash_config.get("target_layer_ids") or config.get("target_layer_ids")
-    if not isinstance(target_layer_ids, list) or not target_layer_ids:
-        raise ValueError("The source checkpoint must define non-empty dflash_config.target_layer_ids.")
-    if any(not isinstance(layer, int) or isinstance(layer, bool) or layer < 0 for layer in target_layer_ids):
-        raise ValueError(f"target_layer_ids must contain non-negative integers; got {target_layer_ids!r}.")
-    if "causal" in dflash_config and not isinstance(dflash_config["causal"], bool):
-        raise ValueError(f"dflash_config.causal must be a boolean; got {dflash_config['causal']!r}.")
-
-    dflash_config.update(
-        {
-            "target_layer_ids": target_layer_ids,
-            "conv_group_size": conv_group_size,
-            "conv_kernel_size": conv_kernel_size,
-            "selector_rank": selector_rank,
-            "selector_top_k": selector_top_k,
-        }
-    )
-    config["architectures"] = [_TARGET_ARCHITECTURE]
-    config["dflash_config"] = dflash_config
-    if "is_causal" not in config and "causal" in dflash_config:
-        config["is_causal"] = bool(dflash_config["causal"])
-    config.pop("num_target_layers", None)
-    quantization_config = config.get("quantization_config")
-    if isinstance(quantization_config, dict):
-        for key in ("ignore", "exclude_modules"):
-            _extend_patterns(quantization_config, key)
-    return config
-
-
-def dflash2_tensor_shapes(config: dict[str, Any]) -> dict[str, tuple[int, ...]]:
-    """Describe the additional tensors required by vLLM's DFlash2 model."""
-    hidden_size = _positive_int(config, "hidden_size")
-    num_layers = _positive_int(config, "num_hidden_layers")
-    vocab_size = _positive_int(config, "vocab_size")
-    dflash_config = config.get("dflash_config") or {}
-    group_size = int(dflash_config["conv_group_size"])
-    kernel_size = int(dflash_config["conv_kernel_size"])
-    selector_rank = int(dflash_config["selector_rank"])
-    if hidden_size % group_size:
-        raise ValueError(f"conv_group_size={group_size} must divide hidden_size={hidden_size}.")
-
-    num_groups = hidden_size // group_size
-    shapes: dict[str, tuple[int, ...]] = {}
-    for layer in range(num_layers):
-        for name in ("attention_conv", "mlp_conv"):
-            prefix = f"layers.{layer}.{name}"
-            shapes[f"{prefix}.base_kernel"] = (2, kernel_size, hidden_size)
-            # ReplicatedLinear stores [output, input]. The output rows flatten
-            # vLLM's (side, tap, group) coefficient layout in that order.
-            shapes[f"{prefix}.kernel_projection.weight"] = (
-                2 * kernel_size * num_groups,
-                hidden_size,
-            )
-    shapes["candidate_selector.predecessor_codebook"] = (vocab_size, selector_rank)
-    shapes["candidate_selector.successor_codebook"] = (vocab_size, selector_rank)
-    shapes["candidate_selector.hidden_projection.weight"] = (selector_rank, hidden_size)
-    return shapes
-
-
-def initialize_dflash2_tensors(config: dict[str, Any]):
-    """Create identity convolutions and a neutral selector in BF16.
-
-    DFlash2's unquantized runtime modules load these tensors in the model dtype.
-    BF16 deliberately matches the Lightning target and keeps the bootstrap
-    memory-representative of a trained Lightning DFlash2 checkpoint.
-    """
-    try:
-        import torch
-    except ImportError as error:
-        raise RuntimeError("Checkpoint conversion requires PyTorch.") from error
-
-    tensors = {}
-    for name, shape in dflash2_tensor_shapes(config).items():
-        tensor = torch.zeros(shape, dtype=torch.bfloat16)
-        if name.endswith(".base_kernel"):
-            # Both the pre-attention/MLP and post-attention/MLP convolutions
-            # pass the current token through unchanged. All older taps and all
-            # dynamic coefficients remain zero.
-            tensor[:, 0, :].fill_(1)
-        tensors[name] = tensor
-    return tensors
-
-
-def convert_checkpoint(
-    source: str,
-    output: str | Path,
-    *,
-    conv_group_size: int = DEFAULT_CONV_GROUP_SIZE,
-    conv_kernel_size: int = DEFAULT_CONV_KERNEL_SIZE,
-    selector_rank: int = DEFAULT_SELECTOR_RANK,
-    selector_top_k: int = DEFAULT_SELECTOR_TOP_K,
-) -> Path:
-    """Convert ``source`` into a new local DFlash2 checkpoint directory."""
-    source_dir, source_label = _resolve_source(source)
-    config_path = source_dir / "config.json"
-    if not config_path.is_file():
-        raise FileNotFoundError(f"Missing source config: {config_path}")
-    source_config = json.loads(config_path.read_text())
-    config = build_dflash2_config(
-        source_config,
-        conv_group_size=conv_group_size,
-        conv_kernel_size=conv_kernel_size,
-        selector_rank=selector_rank,
-        selector_top_k=selector_top_k,
-    )
-
-    output_dir = Path(output).expanduser().resolve()
-    if output_dir.exists():
-        raise FileExistsError(f"Output already exists: {output_dir}")
-    output_dir.parent.mkdir(parents=True, exist_ok=True)
-    with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}-", dir=output_dir.parent) as staging:
-        staging_dir = Path(staging)
-        _write_checkpoint(source_dir, staging_dir, config, source_label)
-        staging_dir.rename(output_dir)
-    return output_dir
-
-
-def _positive_int(config: dict[str, Any], key: str) -> int:
-    value = config.get(key)
-    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
-        raise ValueError(f"config.{key} must be a positive integer; got {value!r}.")
-    return value
-
-
-def _extend_patterns(config: dict[str, Any], key: str) -> None:
-    patterns = config.get(key)
-    if patterns is None:
-        patterns = []
-    if not isinstance(patterns, list) or any(not isinstance(pattern, str) for pattern in patterns):
-        raise ValueError(f"quantization_config.{key} must be a list of strings; got {patterns!r}.")
-    config[key] = [*patterns, *(pattern for pattern in _UNQUANTIZED_MODULE_PATTERNS if pattern not in patterns)]
-
-
-def _resolve_source(source: str) -> tuple[Path, str]:
-    source_path = Path(source).expanduser()
-    if source_path.is_dir():
-        return source_path.resolve(), str(source_path.resolve())
-
-    try:
-        from huggingface_hub import snapshot_download
-    except ImportError as error:
-        raise RuntimeError(
-            "Resolving a Hugging Face model ID requires huggingface_hub. "
-            "Install NeMo's standard dependencies or pass a local checkpoint directory."
-        ) from error
-
-    snapshot = snapshot_download(
-        source,
-        allow_patterns=["*.json", "*.safetensors", "mask_embedding.pt", "LICENSE*", "README*"],
-    )
-    return Path(snapshot), source
-
-
-def _copy_metadata(source_dir: Path, output_dir: Path) -> None:
-    for path in source_dir.glob("LICENSE*"):
-        shutil.copy2(path, output_dir / path.name)
-    mask_embedding = source_dir / "mask_embedding.pt"
-    if mask_embedding.is_file():
-        shutil.copy2(mask_embedding, output_dir / mask_embedding.name)
-
-    source_quant_config = source_dir / "hf_quant_config.json"
-    if source_quant_config.is_file():
-        quant_config = json.loads(source_quant_config.read_text())
-        quantization = quant_config.get("quantization")
-        if isinstance(quantization, dict):
-            _extend_patterns(quantization, "exclude_modules")
-        (output_dir / source_quant_config.name).write_text(json.dumps(quant_config, indent=2) + "\n")
-
-
-def _write_checkpoint(source_dir: Path, output_dir: Path, config: dict[str, Any], source_label: str) -> None:
-    try:
-        from safetensors import safe_open
-        from safetensors.torch import save_file
-    except ImportError as error:
-        raise RuntimeError("Checkpoint conversion requires safetensors.") from error
-
-    weight_files = sorted(source_dir.glob("*.safetensors"))
-    if len(weight_files) != 1:
-        raise ValueError(
-            "This converter expects one safetensors file in the source checkpoint; "
-            f"found {[path.name for path in weight_files]!r}."
-        )
-
-    source_weights = weight_files[0]
-    with safe_open(source_weights, framework="pt", device="cpu") as handle:
-        metadata = handle.metadata()
-        tensors = {name: handle.get_tensor(name) for name in handle.keys()}
-
-    additions = initialize_dflash2_tensors(config)
-    overlap = sorted(tensors.keys() & additions.keys())
-    if overlap:
-        raise ValueError(f"Source checkpoint already contains DFlash2 tensors: {overlap[:5]!r}.")
-    tensors.update(additions)
-
-    save_file(tensors, output_dir / "model.safetensors", metadata=metadata)
-    (output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")
-    manifest = {
-        "source": source_label,
-        "initialization": {
-            "convolutions": "identity",
-            "candidate_selector": "neutral",
-        },
-        "trained_dflash2_parameters": False,
-    }
-    (output_dir / "dflash2_bootstrap.json").write_text(json.dumps(manifest, indent=2) + "\n")
-    (output_dir / "README.md").write_text(
-        "# Lightning DFlash2 bootstrap\n\n"
-        f"This checkpoint was derived from `{source_label}`. The trained DFlash "
-        "backbone is preserved, while the DFlash2 convolutions are initialized as "
-        "identities and its candidate selector is neutral. It is a functional "
-        "DFlash2 runtime artifact, not a DFlash2-trained performance release.\n"
-    )
-    _copy_metadata(source_dir, output_dir)
-
-    # TemporaryDirectory and safetensors default to 0700/0600. The converted
-    # public model is commonly bind-mounted into a root-squashed inference
-    # container, so normalize it to read-only model-artifact permissions.
-    output_dir.chmod(0o755)
-    for path in output_dir.iterdir():
-        if path.is_file():
-            path.chmod(0o644)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("source", help="Local DFlash checkpoint directory or Hugging Face model ID")
-    parser.add_argument("output", help="New local output directory")
-    parser.add_argument("--conv-group-size", type=int, default=DEFAULT_CONV_GROUP_SIZE)
-    parser.add_argument("--conv-kernel-size", type=int, default=DEFAULT_CONV_KERNEL_SIZE)
-    parser.add_argument("--selector-rank", type=int, default=DEFAULT_SELECTOR_RANK)
-    parser.add_argument("--selector-top-k", type=int, default=DEFAULT_SELECTOR_TOP_K)
-    args = parser.parse_args()
-    output = convert_checkpoint(
-        args.source,
-        args.output,
-        conv_group_size=args.conv_group_size,
-        conv_kernel_size=args.conv_kernel_size,
-        selector_rank=args.selector_rank,
-        selector_top_k=args.selector_top_k,
-    )
-    print(f"Created DFlash2 checkpoint: {output}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/collections/speechlm2/test_dflash2_checkpoint.py b/tests/collections/speechlm2/test_dflash2_checkpoint.py
deleted file mode 100644
index 6de12b5b7a42..000000000000
--- a/tests/collections/speechlm2/test_dflash2_checkpoint.py
+++ /dev/null
@@ -1,281 +0,0 @@
-# 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.
-
-import json
-
-import pytest
-
-from scripts.speechlm2.convert_dflash_to_dflash2 import (
-    build_dflash2_config,
-    convert_checkpoint,
-    dflash2_tensor_shapes,
-    initialize_dflash2_tensors,
-)
-
-
-@pytest.fixture
-def lightning_dflash_config():
-    return {
-        "architectures": ["DFlashDraftModel"],
-        "hidden_size": 2688,
-        "num_hidden_layers": 6,
-        "num_target_layers": 52,
-        "vocab_size": 131072,
-        "dflash_config": {
-            "causal": False,
-            "mask_token_id": 990,
-            "target_layer_ids": [1, 5, 19, 29, 41, 51],
-        },
-        "quantization_config": {
-            "ignore": ["*embed_tokens*", "*self_attn*"],
-            "exclude_modules": ["*embed_tokens*", "*self_attn*"],
-        },
-    }
-
-
-@pytest.fixture
-def tiny_dflash_checkpoint(tmp_path, lightning_dflash_config):
-    torch = pytest.importorskip("torch")
-    safetensors = pytest.importorskip("safetensors.torch")
-    source = tmp_path / "source"
-    source.mkdir()
-    config = {
-        **lightning_dflash_config,
-        "hidden_size": 8,
-        "num_hidden_layers": 1,
-        "vocab_size": 32,
-        "dflash_config": {
-            **lightning_dflash_config["dflash_config"],
-            "target_layer_ids": [1],
-        },
-    }
-    (source / "config.json").write_text(json.dumps(config))
-    (source / "hf_quant_config.json").write_text(
-        json.dumps({"quantization": {"exclude_modules": ["*embed_tokens*", "*self_attn*"]}})
-    )
-    (source / "mask_embedding.pt").write_bytes(b"test mask embedding")
-    safetensors.save_file({"norm.weight": torch.arange(8, dtype=torch.bfloat16)}, source / "model.safetensors")
-    return source
-
-
-def test_build_config_is_lightning_compatible(lightning_dflash_config):
-    converted = build_dflash2_config(lightning_dflash_config)
-
-    assert converted["architectures"] == ["DFlash2DraftModel"]
-    assert "num_target_layers" not in converted
-    assert converted["is_causal"] is False
-    assert converted["dflash_config"] == {
-        "causal": False,
-        "mask_token_id": 990,
-        "target_layer_ids": [1, 5, 19, 29, 41, 51],
-        "conv_group_size": 16,
-        "conv_kernel_size": 2,
-        "selector_rank": 256,
-        "selector_top_k": 16,
-    }
-    expected_exclusions = [
-        "*embed_tokens*",
-        "*self_attn*",
-        "*attention_conv*",
-        "*mlp_conv*",
-        "*candidate_selector*",
-    ]
-    assert converted["quantization_config"]["ignore"] == expected_exclusions
-    assert converted["quantization_config"]["exclude_modules"] == expected_exclusions
-
-
-def test_build_config_preserves_layer_type_causality(lightning_dflash_config):
-    del lightning_dflash_config["dflash_config"]["causal"]
-    lightning_dflash_config["layer_types"] = ["sliding_attention", "full_attention"]
-
-    converted = build_dflash2_config(lightning_dflash_config)
-
-    assert "is_causal" not in converted
-    assert converted["layer_types"] == ["sliding_attention", "full_attention"]
-
-
-def test_build_config_normalizes_top_level_target_layers(lightning_dflash_config):
-    target_layer_ids = lightning_dflash_config["dflash_config"].pop("target_layer_ids")
-    lightning_dflash_config["target_layer_ids"] = target_layer_ids
-
-    converted = build_dflash2_config(lightning_dflash_config)
-
-    assert converted["dflash_config"]["target_layer_ids"] == target_layer_ids
-
-
-def test_tensor_shapes_match_vllm_dflash2(lightning_dflash_config):
-    shapes = dflash2_tensor_shapes(build_dflash2_config(lightning_dflash_config))
-
-    assert len(shapes) == 27
-    assert shapes["layers.0.attention_conv.base_kernel"] == (2, 2, 2688)
-    assert shapes["layers.5.mlp_conv.kernel_projection.weight"] == (672, 2688)
-    assert shapes["candidate_selector.predecessor_codebook"] == (131072, 256)
-    assert shapes["candidate_selector.successor_codebook"] == (131072, 256)
-    assert shapes["candidate_selector.hidden_projection.weight"] == (256, 2688)
-
-
-def test_tensor_shapes_match_installed_vllm_modules(monkeypatch, lightning_dflash_config):
-    torch = pytest.importorskip("torch")
-    linear = pytest.importorskip("vllm.model_executor.layers.linear")
-    parameter = pytest.importorskip("vllm.model_executor.parameter")
-    dflash2 = pytest.importorskip("vllm.model_executor.models.qwen3_dflash2")
-    vllm_config = pytest.importorskip("vllm.config")
-
-    for module in (linear, parameter):
-        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
-        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
-
-    config = {
-        **lightning_dflash_config,
-        "hidden_size": 8,
-        "num_hidden_layers": 1,
-        "vocab_size": 32,
-    }
-    config = build_dflash2_config(config, conv_group_size=4, selector_rank=2, selector_top_k=4)
-    with vllm_config.set_current_vllm_config(vllm_config.VllmConfig()):
-        conv = dflash2.DFlashGroupedConv(
-            hidden_size=8,
-            taps=2,
-            group_size=4,
-            block_size=3,
-            params_dtype=torch.bfloat16,
-            prefix="attention_conv",
-        )
-        selector = dflash2.CandidateSelector(
-            hidden_size=8,
-            vocab_size=32,
-            rank=2,
-            top_k=4,
-            params_dtype=torch.bfloat16,
-            prefix="candidate_selector",
-        )
-
-    actual = {
-        **{f"layers.0.attention_conv.{name}": tuple(param.shape) for name, param in conv.named_parameters()},
-        **{f"candidate_selector.{name}": tuple(param.shape) for name, param in selector.named_parameters()},
-    }
-    expected = dflash2_tensor_shapes(config)
-    covered_expected = {
-        name for name in expected if name.startswith(("layers.0.attention_conv.", "candidate_selector."))
-    }
-    assert set(actual) == covered_expected
-    for name, shape in actual.items():
-        assert shape == expected[name]
-
-
-def test_identity_convolutions_and_neutral_selector(lightning_dflash_config):
-    torch = pytest.importorskip("torch")
-    tensors = initialize_dflash2_tensors(build_dflash2_config(lightning_dflash_config))
-
-    base = tensors["layers.0.attention_conv.base_kernel"]
-    torch.testing.assert_close(base[:, 0], torch.ones_like(base[:, 0]))
-    torch.testing.assert_close(base[:, 1], torch.zeros_like(base[:, 1]))
-    assert not tensors["layers.0.attention_conv.kernel_projection.weight"].count_nonzero()
-    assert not tensors["candidate_selector.predecessor_codebook"].count_nonzero()
-    assert not tensors["candidate_selector.successor_codebook"].count_nonzero()
-    assert not tensors["candidate_selector.hidden_projection.weight"].count_nonzero()
-
-
-def test_convert_checkpoint_writes_loadable_artifact(tmp_path, tiny_dflash_checkpoint):
-    torch = pytest.importorskip("torch")
-    safetensors = pytest.importorskip("safetensors.torch")
-    original = safetensors.load_file(tiny_dflash_checkpoint / "model.safetensors")["norm.weight"]
-
-    output = convert_checkpoint(
-        str(tiny_dflash_checkpoint),
-        tmp_path / "output",
-        conv_group_size=4,
-        selector_rank=2,
-        selector_top_k=4,
-    )
-
-    output_config = json.loads((output / "config.json").read_text())
-    output_weights = safetensors.load_file(output / "model.safetensors")
-    manifest = json.loads((output / "dflash2_bootstrap.json").read_text())
-    hf_quant_config = json.loads((output / "hf_quant_config.json").read_text())
-    assert output_config["architectures"] == ["DFlash2DraftModel"]
-    torch.testing.assert_close(output_weights["norm.weight"], original)
-    assert output_weights["layers.0.attention_conv.kernel_projection.weight"].shape == (8, 8)
-    assert manifest["trained_dflash2_parameters"] is False
-    assert (output / "hf_quant_config.json").is_file()
-    assert (output / "mask_embedding.pt").read_bytes() == b"test mask embedding"
-    assert hf_quant_config["quantization"]["exclude_modules"] == [
-        "*embed_tokens*",
-        "*self_attn*",
-        "*attention_conv*",
-        "*mlp_conv*",
-        "*candidate_selector*",
-    ]
-    assert output.stat().st_mode & 0o777 == 0o755
-    assert (output / "model.safetensors").stat().st_mode & 0o777 == 0o644
-
-
-@pytest.mark.parametrize(
-    ("kwargs", "match"),
-    [
-        ({"conv_group_size": 17}, "must divide hidden_size"),
-        ({"selector_top_k": 131073}, "cannot exceed vocab_size"),
-        ({"conv_kernel_size": 0}, "must be a positive integer"),
-    ],
-)
-def test_rejects_incompatible_dimensions(lightning_dflash_config, kwargs, match):
-    with pytest.raises(ValueError, match=match):
-        build_dflash2_config(lightning_dflash_config, **kwargs)
-
-
-def test_rejects_non_dflash_source(lightning_dflash_config):
-    lightning_dflash_config["architectures"] = ["DFlash2DraftModel"]
-
-    with pytest.raises(ValueError, match="must declare DFlashDraftModel"):
-        build_dflash2_config(lightning_dflash_config)
-
-
-def test_rejects_existing_output(tmp_path, tiny_dflash_checkpoint):
-    output = tmp_path / "output"
-    output.mkdir()
-
-    with pytest.raises(FileExistsError, match="Output already exists"):
-        convert_checkpoint(str(tiny_dflash_checkpoint), output, conv_group_size=4, selector_rank=2)
-
-
-def test_rejects_multiple_weight_files(tmp_path, tiny_dflash_checkpoint):
-    (tiny_dflash_checkpoint / "second.safetensors").write_bytes(b"not read")
-
-    with pytest.raises(ValueError, match="expects one safetensors file"):
-        convert_checkpoint(
-            str(tiny_dflash_checkpoint),
-            tmp_path / "output",
-            conv_group_size=4,
-            selector_rank=2,
-        )
-
-
-def test_rejects_existing_dflash2_tensors(tmp_path, tiny_dflash_checkpoint):
-    torch = pytest.importorskip("torch")
-    safetensors = pytest.importorskip("safetensors.torch")
-    safetensors.save_file(
-        {
-            "norm.weight": torch.arange(8, dtype=torch.bfloat16),
-            "layers.0.attention_conv.base_kernel": torch.zeros((2, 2, 8), dtype=torch.bfloat16),
-        },
-        tiny_dflash_checkpoint / "model.safetensors",
-    )
-
-    with pytest.raises(ValueError, match="already contains DFlash2 tensors"):
-        convert_checkpoint(
-            str(tiny_dflash_checkpoint),
-            tmp_path / "output",
-            conv_group_size=4,
-            selector_rank=2,
-        )

From 61b7176277fee36f89d5eeacdf16ee119f155a9d Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Mon, 24 Aug 2026 17:10:36 -0700
Subject: [PATCH 17/20] Pin vLLM DFlash2 runtime commit

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 docs/source/speechlm2/vllm_dflash.rst | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
index bd1bcb2860e3..06b9c12b7315 100644
--- a/docs/source/speechlm2/vllm_dflash.rst
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -41,13 +41,15 @@ model. Its checkpoint must be trained or fine-tuned separately for the target
 language backbone; NeMo's vLLM inference plugin does not create or convert
 DFlash2 weights.
 
-At the time of writing, DFlash2 requires the vLLM implementation from pull
-request 52816. It uses the same ``method`` value as DFlash; vLLM selects the
-DFlash2 runtime from the trained draft checkpoint's architecture:
+At the time of writing, DFlash2 requires vLLM commit
+``3406ec1dae9916f920b90f0dbf90dcf54923d042`` from pull request 52816. The
+immutable commit pin keeps the DFlash2 runtime reproducible. DFlash2 uses the
+same ``method`` value as DFlash; vLLM selects it from the trained draft
+checkpoint's architecture:
 
 .. code-block:: bash
 
-   pip install -U "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/52816/head"
+   pip install -U "vllm @ git+https://github.com/vllm-project/vllm.git@3406ec1dae9916f920b90f0dbf90dcf54923d042"
 
    vllm serve /path/to/vllm-ready-speechlm-checkpoint \
      --trust-remote-code \

From 71c2ef18ddff1d93435a8704ddd69ab51ba2204e Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 22:44:18 -0700
Subject: [PATCH 18/20] fix(speechlm2): load embedded PE inference checkpoints

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 examples/speechlm2/to_hf.py                   |   2 +-
 .../asr/modules/parallel_expert_encoder.py    |  22 ++--
 nemo/collections/speechlm2/vllm/salm/audio.py | 114 ++++++++++++++----
 nemo/collections/speechlm2/vllm/salm/model.py |   6 +-
 .../asr/test_parallel_expert_encoder.py       |  35 ++++++
 tests/collections/speechlm2/test_to_hf.py     |   2 +-
 .../collections/speechlm2/test_vllm_plugin.py |  54 +++++++++
 7 files changed, 196 insertions(+), 39 deletions(-)

diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
index a5d6159ee8dc..55996171fc92 100644
--- a/examples/speechlm2/to_hf.py
+++ b/examples/speechlm2/to_hf.py
@@ -263,7 +263,7 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     if existing:
         LOG.info("Overwriting existing files in %s: %s", output_dir, existing)
     tokenizer_src = model_cfg.get("tokenizer_path") or pretrained_llm
-    tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True)
+    tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True, fix_mistral_regex=True)
     if audio_token not in tok.get_vocab():
         tok.add_special_tokens({"additional_special_tokens": [audio_token]})
     audio_token_id = tok.get_vocab().get(audio_token)
diff --git a/nemo/collections/asr/modules/parallel_expert_encoder.py b/nemo/collections/asr/modules/parallel_expert_encoder.py
index 73e1c488c85d..e626e9dfa751 100644
--- a/nemo/collections/asr/modules/parallel_expert_encoder.py
+++ b/nemo/collections/asr/modules/parallel_expert_encoder.py
@@ -14,10 +14,10 @@
 
 """Parallel Expert Speech Encoder.
 
-Runs a Sortformer speaker-diarization expert and an ASR Conformer encoder on the
+Runs a Sortformer speaker-diarization expert and an ASR encoder on the
 same mel input, then fuses their outputs (LayerNorm + sinusoidal speaker-kernel +
 ADD). Expects un-normalised mels; the ASR branch re-applies ``normalize_batch``
-internally. I/O matches :class:`ConformerEncoder` (drop-in). Only self-contained PE
+internally. I/O matches the supported ASR encoders (drop-in). Only self-contained PE
 bundles (inline ``asr_encoder_cfg`` + ``diarization_model_cfg`` in
 ``model_config.yaml``) are supported.
 """
@@ -38,6 +38,7 @@
 from tqdm import tqdm
 
 from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder
 from nemo.collections.asr.parts.preprocessing.features import normalize_batch
 from nemo.core.classes import ModelPT
 from nemo.core.classes.common import PretrainedModelInfo
@@ -284,16 +285,17 @@ def save_to_nemo(
 
 @experimental
 class ParallelExpertEncoder(nn.Module):
-    """Sortformer-diarizer + ASR Conformer encoder; I/O identical to :class:`ConformerEncoder`.
+    """Sortformer-diarizer plus a supported ASR encoder with a shared I/O contract.
 
     Reconstructed from inline configs in the PE bundle's ``model_config.yaml``.
 
     Args:
-        asr_encoder_cfg (DictConfig): Inline config for the ASR-side :class:`ConformerEncoder`.
+        asr_encoder_cfg (DictConfig): Inline config for an ASR-side :class:`ConformerEncoder`
+            or :class:`TransformerEncoder`.
         diarization_model_cfg (DictConfig): Inline config for the :class:`SortformerEncLabelModel`.
         asr_normalize_type (str, optional): Normalization replayed on the ASR branch. Defaults to ``per_feature``.
         freeze_diar (bool): Freeze the Sortformer parameters. Defaults to ``True``.
-        freeze_asr (bool): Freeze the wrapped ASR ConformerEncoder. Defaults to ``False``.
+        freeze_asr (bool): Freeze the wrapped ASR encoder. Defaults to ``False``.
         online_inference_length (int): Online-inference window in encoder output frames
             (default ``500`` ~= 40s); ``<= 0`` disables it.
         chunk_left_context (int): Left context (output frames) per online window, shared by
@@ -332,10 +334,10 @@ def __init__(
             )
 
         self.asr_encoder = ConformerEncoder.from_config_dict(_clone_config(asr_encoder_cfg))
-        if not isinstance(self.asr_encoder, ConformerEncoder):
+        if not isinstance(self.asr_encoder, (ConformerEncoder, TransformerEncoder)):
             raise TypeError(
-                f"Expected `asr_encoder_cfg._target_` to instantiate a "
-                f"ConformerEncoder, got {type(self.asr_encoder).__name__} instead."
+                "Expected `asr_encoder_cfg._target_` to instantiate a ConformerEncoder "
+                f"or TransformerEncoder, got {type(self.asr_encoder).__name__} instead."
             )
         self.asr_normalize_type = asr_normalize_type or 'per_feature'
         self._feat_in = self.asr_encoder._feat_in
@@ -406,7 +408,7 @@ def train(self, mode: bool = True) -> "ParallelExpertEncoder":
             self.asr_encoder.eval()
         return self
 
-    # ConformerEncoder-compatible properties (drop-in for SALM perception).
+    # ASR-encoder-compatible properties (drop-in for SALM perception).
     @property
     def d_model(self) -> int:
         return self.asr_d_model
@@ -486,7 +488,7 @@ def _fuse_diar_and_asr(self, asr_encoded: torch.Tensor, spk_targets: torch.Tenso
 
         return fused.transpose(1, 2)  # (B, D, T)
 
-    # Forward — identical signature to ConformerEncoder.forward
+    # Forward — identical signature across the supported ASR encoders.
     def forward(
         self,
         audio_signal,
diff --git a/nemo/collections/speechlm2/vllm/salm/audio.py b/nemo/collections/speechlm2/vllm/salm/audio.py
index d55f3b827308..036bf476aeac 100644
--- a/nemo/collections/speechlm2/vllm/salm/audio.py
+++ b/nemo/collections/speechlm2/vllm/salm/audio.py
@@ -108,7 +108,64 @@ def _load_nemo_perception(perception_cfg: dict) -> nn.Module:
     return perception
 
 
-def _maybe_mount_pe_encoder(perception: nn.Module, pe_encoder_path: str | None) -> bool:
+def _build_pe_encoder_from_config(pe_encoder_config: Mapping[str, object]) -> nn.Module:
+    """Reconstruct a PE encoder whose config and weights are embedded in an HF export."""
+    from omegaconf import OmegaConf
+
+    from nemo.collections.asr.modules.parallel_expert_encoder import (
+        ParallelExpertEncoder,
+    )
+
+    if hasattr(pe_encoder_config, "to_dict"):
+        pe_encoder_config = pe_encoder_config.to_dict()
+    if not isinstance(pe_encoder_config, Mapping):
+        raise TypeError(
+            "pe_encoder_config must be a mapping containing inline ASR and diarization configs; "
+            f"got {type(pe_encoder_config).__name__}."
+        )
+
+    cfg = OmegaConf.create(dict(pe_encoder_config))
+    return ParallelExpertEncoder(
+        asr_encoder_cfg=cfg.get("asr_encoder_cfg"),
+        diarization_model_cfg=cfg.get("diarization_model_cfg"),
+        asr_normalize_type=cfg.get("asr_normalize_type"),
+        freeze_diar=cfg.get("freeze_diar", True),
+        freeze_asr=cfg.get("freeze_asr", False),
+        online_inference_length=cfg.get("online_inference_length", 500),
+        chunk_left_context=cfg.get("chunk_left_context", 50),
+        chunk_right_context=cfg.get("chunk_right_context", 50),
+        diar_fifo_len=cfg.get("diar_fifo_len", 40),
+        diar_spkcache_update_period=cfg.get("diar_spkcache_update_period", 300),
+        diar_spkcache_len=cfg.get("diar_spkcache_len", 188),
+    )
+
+
+def _attach_pe_encoder(perception: nn.Module, pe_encoder: nn.Module, source: str) -> None:
+    """Replace the direct perception encoder while preserving placement and preprocessing."""
+    if not hasattr(perception, "encoder"):
+        raise RuntimeError(f"{source} is set but perception has no encoder attribute to replace.")
+
+    existing_d_model = int(getattr(perception.encoder, "d_model", -1))
+    if existing_d_model > 0 and int(pe_encoder.d_model) != existing_d_model:
+        raise ValueError(
+            f"ParallelExpertEncoder d_model={pe_encoder.d_model} does not match the existing "
+            f"perception encoder d_model={existing_d_model}."
+        )
+
+    ref_param = next(perception.encoder.parameters(), None)
+    if ref_param is not None:
+        pe_encoder = pe_encoder.to(device=ref_param.device, dtype=ref_param.dtype)
+
+    try:
+        perception.preprocessor.featurizer.normalize = None
+    except AttributeError:
+        pass
+
+    perception.encoder = pe_encoder
+    perception.eval()
+
+
+def _mount_pe_encoder_from_path(perception: nn.Module, pe_encoder_path: str | None) -> bool:
     """Replace ``perception.encoder`` with a ParallelExpertEncoder bundle so PE-trained
     checkpoints (nested ``asr_encoder.*`` / ``diarization_model.*`` weights) load correctly.
 
@@ -135,45 +192,50 @@ def _maybe_mount_pe_encoder(perception: nn.Module, pe_encoder_path: str | None)
     """
     if pe_encoder_path in (None, "", False):
         return False
-    if not hasattr(perception, "encoder"):
-        raise RuntimeError("pe_encoder_path is set but perception has no `encoder` attribute to replace.")
+    if not isinstance(pe_encoder_path, str):
+        raise TypeError(f"pe_encoder_path must be a string, got {type(pe_encoder_path).__name__}.")
 
     from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoderPT
 
     # Only fail-fast for a *local* ``.nemo`` file that is not a PE bundle. A
     # non-local reference (HF repo id / NGC alias) is resolved offline from the
     # HuggingFace cache by load_from_nemo -> from_pretrained, so do not reject it.
-    is_local_nemo_file = (
-        isinstance(pe_encoder_path, str) and pe_encoder_path.endswith(".nemo") and os.path.isfile(pe_encoder_path)
-    )
+    is_local_nemo_file = pe_encoder_path.endswith(".nemo") and os.path.isfile(pe_encoder_path)
     if is_local_nemo_file and not ParallelExpertEncoderPT.is_pe_nemo(pe_encoder_path):
         raise ValueError(f"pe_encoder_path={pe_encoder_path!r} is not a ParallelExpertEncoderPT .nemo bundle.")
 
     pe_encoder = ParallelExpertEncoderPT.load_from_nemo(pe_encoder_path, map_location="cpu", strict=True)
+    _attach_pe_encoder(perception, pe_encoder, "pe_encoder_path")
+    return True
 
-    existing_d_model = int(getattr(perception.encoder, "d_model", -1))
-    if existing_d_model > 0 and int(pe_encoder.d_model) != existing_d_model:
-        raise ValueError(
-            f"ParallelExpertEncoder d_model={pe_encoder.d_model} does not match the existing "
-            f"perception encoder d_model={existing_d_model}."
-        )
 
-    # load_from_nemo restores onto CPU; copy the replaced encoder's device/dtype to avoid CPU/dtype mismatches.
-    ref_param = next(perception.encoder.parameters(), None)
-    if ref_param is not None:
-        pe_encoder = pe_encoder.to(device=ref_param.device, dtype=ref_param.dtype)
+def _maybe_mount_pe_encoder(
+    perception: nn.Module,
+    pe_encoder_path: str | None,
+    pe_encoder_config: Mapping[str, object] | None = None,
+) -> bool:
+    """Mount the PE encoder represented by an exported SpeechLM checkpoint.
 
-    # PE encoder consumes un-normalised mels and replays ASR norm internally, so disable preprocessor norm.
-    try:
-        perception.preprocessor.featurizer.normalize = None
-    except AttributeError:
-        # Preprocessor/featurizer layout varies across backends; if the attribute is
-        # absent there is no outer normalization to disable, so skipping is correct.
-        pass
+    Self-contained HF exports store the constructor inputs in pe_encoder_config
+    and the nested encoder weights in their own safetensors file. That embedded
+    form takes precedence over pe_encoder_path so serving does not depend on a
+    path from the training machine. Legacy path-only checkpoints retain their
+    existing local-bundle and pretrained-model behavior.
 
-    perception.encoder = pe_encoder
-    perception.eval()
-    return True
+    Args:
+        perception: Perception module whose direct encoder is replaced.
+        pe_encoder_path: Legacy local bundle path or pretrained model identifier.
+        pe_encoder_config: Optional self-contained PE constructor config.
+
+    Returns:
+        True when a PE encoder was mounted, otherwise False.
+    """
+    if pe_encoder_config:
+        pe_encoder = _build_pe_encoder_from_config(pe_encoder_config)
+        _attach_pe_encoder(perception, pe_encoder, "pe_encoder_config")
+        return True
+
+    return _mount_pe_encoder_from_path(perception, pe_encoder_path)
 
 
 def _pad_to_vocab_size(tensor: torch.Tensor, target_vocab: int) -> torch.Tensor:
diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py
index d73ce1fa8bf9..d2979df36fc7 100644
--- a/nemo/collections/speechlm2/vllm/salm/model.py
+++ b/nemo/collections/speechlm2/vllm/salm/model.py
@@ -109,7 +109,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         with self._mark_tower_model(vllm_config, {"audio"}):
             self.perception = _load_nemo_perception(config.perception)
-            _maybe_mount_pe_encoder(self.perception, getattr(config, "pe_encoder_path", None))
+            _maybe_mount_pe_encoder(
+                self.perception,
+                getattr(config, "pe_encoder_path", None),
+                getattr(config, "pe_encoder_config", None),
+            )
 
         self._uses_pe_encoder = isinstance(getattr(self.perception, "encoder", None), ParallelExpertEncoder)
 
diff --git a/tests/collections/asr/test_parallel_expert_encoder.py b/tests/collections/asr/test_parallel_expert_encoder.py
index b9c099cc6661..d7a57ec94ddc 100644
--- a/tests/collections/asr/test_parallel_expert_encoder.py
+++ b/tests/collections/asr/test_parallel_expert_encoder.py
@@ -30,6 +30,7 @@
     _default_dtype,
     _disable_dist_feature_sync,
 )
+from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder
 
 # ``@experimental`` wraps the class in a wrapt proxy, so ``__new__`` (used to build
 # bare instances that skip the heavy real ``__init__``) must target the underlying
@@ -347,6 +348,29 @@ def toy_asr_encoder_cfg() -> DictConfig:
     )
 
 
+def toy_transformer_asr_encoder_cfg() -> DictConfig:
+    """Tiny production-style TransformerEncoder config for the PE ASR branch."""
+    return DictConfig(
+        {
+            '_target_': 'nemo.collections.asr.modules.transformer_encoder.TransformerEncoder',
+            'feat_in': _MEL_FEATURES,
+            'd_model': 64,
+            'n_heads': 4,
+            'n_layers': 1,
+            'drop_rate': 0.0,
+            'qkv_bias': False,
+            'qk_norm': True,
+            'ff_expansion': 2,
+            'pre_block_norm': True,
+            'subsampling_factor': _SUBSAMPLING_FACTOR,
+            'attn_mode': 'full',
+            'self_attention_model': 'rope',
+            'rope_base': 10000.0,
+            'rotary_fraction': 0.5,
+        }
+    )
+
+
 def toy_diarization_model_cfg() -> DictConfig:
     """Tiny SortformerEncLabelModel config the PE encoder mounts as its diar branch."""
     model_defaults = {'fc_d_model': _DIAR_FC_D_MODEL, 'tf_d_model': _DIAR_TF_D_MODEL}
@@ -462,6 +486,17 @@ def test_pe_encoder_builds_and_wires_both_real_encoders():
     assert any(p.requires_grad for p in enc.asr_encoder.parameters())
 
 
+@pytest.mark.unit
+def test_pe_encoder_accepts_transformer_asr_branch():
+    enc = build_toy_pe_encoder(asr_encoder_cfg=toy_transformer_asr_encoder_cfg())
+
+    assert isinstance(enc.asr_encoder, TransformerEncoder)
+    assert enc.d_model == 64
+    assert enc.subsampling_factor == _SUBSAMPLING_FACTOR
+    assert enc.pre_encode is enc.asr_encoder.pre_encode
+    assert enc.diar_kernel.shape == (_N_SPK, 64)
+
+
 @pytest.mark.unit
 @pytest.mark.parametrize(
     "high_resolution, requested_diar_subsampling_factor, expected_asr_aligned_factor",
diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py
index 5d9f0fbf5eaa..d8d2309188e9 100644
--- a/tests/collections/speechlm2/test_to_hf.py
+++ b/tests/collections/speechlm2/test_to_hf.py
@@ -445,7 +445,7 @@ def test_prepare_for_vllm_uses_training_tokenizer_path(tmp_path):
             },
         )
 
-    load_tokenizer.assert_called_once_with("custom-tokenizer", trust_remote_code=True)
+    load_tokenizer.assert_called_once_with("custom-tokenizer", trust_remote_code=True, fix_mistral_regex=True)
 
 
 def test_prepare_for_vllm_tokenizer_config_normalized(tmp_path):
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 8f491239edd1..a9481487ad8b 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -612,6 +612,60 @@ def test_call_hf_processor_emits_true_audio_lengths(self):
         assert result["audio_signal"][0].shape[-1] == 12345
         assert torch.equal(result["audio_signal_length"], torch.tensor([12345]))
 
+    def test_embedded_pe_config_takes_precedence_over_training_path(self, monkeypatch):
+        from unittest.mock import Mock
+
+        import torch
+        from torch import nn
+
+        import nemo.collections.asr.modules.parallel_expert_encoder as pe_module
+        from nemo.collections.speechlm2.vllm.salm.audio import _maybe_mount_pe_encoder
+
+        class _Encoder(nn.Module):
+            d_model = 4
+
+            def __init__(self):
+                super().__init__()
+                self.weight = nn.Parameter(torch.ones(1))
+
+        class _Perception(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.encoder = _Encoder().to(torch.float64)
+                self.preprocessor = SimpleNamespace(featurizer=SimpleNamespace(normalize="per_feature"))
+
+        perception = _Perception()
+        embedded_encoder = _Encoder()
+        built = {}
+
+        def build_encoder(**kwargs):
+            built.update(kwargs)
+            return embedded_encoder
+
+        load_from_nemo = Mock(side_effect=AssertionError("embedded config must avoid external path"))
+        monkeypatch.setattr(pe_module.ParallelExpertEncoderPT, "load_from_nemo", load_from_nemo)
+        monkeypatch.setattr(pe_module, "ParallelExpertEncoder", build_encoder)
+
+        mounted = _maybe_mount_pe_encoder(
+            perception,
+            "/training/only/encoder.nemo",
+            {
+                "target": "nemo.collections.asr.modules.parallel_expert_encoder.ParallelExpertEncoderPT",
+                "asr_encoder_cfg": {"kind": "transformer"},
+                "diarization_model_cfg": {"kind": "sortformer"},
+                "asr_normalize_type": "per_feature",
+            },
+        )
+
+        assert mounted
+        assert perception.encoder is embedded_encoder
+        assert perception.encoder.weight.dtype == torch.float64
+        assert perception.preprocessor.featurizer.normalize is None
+        assert not perception.training
+        assert built["asr_encoder_cfg"]["kind"] == "transformer"
+        assert built["diarization_model_cfg"]["kind"] == "sortformer"
+        load_from_nemo.assert_not_called()
+
     def test_perception_forward(self):
         """A small NeMo perception module should encode dummy audio to embeddings."""
         import torch

From c942cd77d1a91146672336df4155a6876b92ca52 Mon Sep 17 00:00:00 2001
From: slyne deng <slyned@nvidia.com>
Date: Wed, 26 Aug 2026 14:45:27 -0700
Subject: [PATCH 19/20] fix(speechlm2): accept Automodel DFlash drafts

Signed-off-by: slyne deng <slyned@nvidia.com>
---
 docs/source/speechlm2/vllm_dflash.rst         |  4 +++
 .../speechlm2/vllm/salm/__init__.py           |  7 +++++
 .../collections/speechlm2/test_vllm_plugin.py | 29 +++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
index 06b9c12b7315..15bf19a78bd0 100644
--- a/docs/source/speechlm2/vllm_dflash.rst
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -32,6 +32,10 @@ The target SpeechLM checkpoint must use
 ``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`` as its language backbone.
 The draft checkpoint provides the auxiliary target-layer selection and mask
 token configuration consumed by vLLM; no draft weights are bundled with NeMo.
+Automodel-trained drafts may retain ``Qwen3DFlashDraftModel`` as their
+architecture so they can be reopened for training. The NeMo plugin registers
+that name as an alias of vLLM's native ``DFlashDraftModel`` implementation;
+serve the original checkpoint without rewriting its ``config.json``.
 
 DFlash2 inference
 -----------------
diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index 40943c3d8e1b..add8c2916f31 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -159,5 +159,12 @@ def register():
         "NeMoSpeechLMForConditionalGeneration",
         f"{_PKG}.model:NeMoSpeechLMForConditionalGeneration",
     )
+    supported_archs = ModelRegistry.get_supported_archs()
+    if "DFlashDraftModel" in supported_archs:
+        native_dflash_model = ModelRegistry.models["DFlashDraftModel"]
+        native_dflash_model_ref = f"{native_dflash_model.module_name}:{native_dflash_model.class_name}"
+        for automodel_arch in ("Qwen3DFlashDraftModel", "DFlashQwen3DFlashDraftModel"):
+            if automodel_arch not in supported_archs:
+                ModelRegistry.register_model(automodel_arch, native_dflash_model_ref)
 
     _patch_vllm_for_nemo_speechlm_mtp()
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index a9481487ad8b..4a7bc2760e20 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -798,6 +798,35 @@ def test_register_does_not_load_backbone_config(self, monkeypatch):
 class TestDFlashPlugin:
     """Tests for the target-model contract required by DFlash and DFlash2."""
 
+    def test_registers_automodel_dflash_architecture_alias(self, monkeypatch):
+        """An untouched Automodel draft config should resolve to vLLM's native DFlash model."""
+        from transformers import AutoConfig
+        from vllm.model_executor.models.registry import ModelRegistry
+        from vllm.transformers_utils.configs.eagle import EAGLEConfig
+
+        if "DFlashDraftModel" not in ModelRegistry.get_supported_archs():
+            pytest.skip("installed vLLM does not provide native DFlash support")
+
+        monkeypatch.setattr(
+            AutoConfig,
+            "from_pretrained",
+            lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()),
+        )
+
+        register()
+
+        native_model = ModelRegistry.models["DFlashDraftModel"]
+        automodel_config = AutoConfig.for_model("qwen3", architectures=["Qwen3DFlashDraftModel"])
+        runtime_config = EAGLEConfig(automodel_config, method="dflash", model_type="eagle")
+        assert runtime_config.architectures == ["DFlashQwen3DFlashDraftModel"]
+
+        for alias in ("Qwen3DFlashDraftModel", *runtime_config.architectures):
+            alias_model = ModelRegistry.models[alias]
+            assert (alias_model.module_name, alias_model.class_name) == (
+                native_model.module_name,
+                native_model.class_name,
+            )
+
     def test_model_advertises_eagle3_support(self):
         from vllm.model_executor.models.interfaces import supports_eagle3
 

From e095c3c4fbf02ff1130d69e8d833670296948fa9 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Fri, 4 Sep 2026 10:05:21 -0700
Subject: [PATCH 20/20] feat(speechlm2): route Automodel DFlash2 through vLLM

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 docs/source/speechlm2/vllm_dflash.rst         | 14 ++++-
 .../speechlm2/vllm/salm/__init__.py           | 41 ++++++++++--
 nemo/collections/speechlm2/vllm/salm/audio.py |  4 +-
 .../collections/speechlm2/test_vllm_plugin.py | 63 +++++++++++++++++++
 4 files changed, 112 insertions(+), 10 deletions(-)

diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst
index 15bf19a78bd0..bf56081ce9ab 100644
--- a/docs/source/speechlm2/vllm_dflash.rst
+++ b/docs/source/speechlm2/vllm_dflash.rst
@@ -63,9 +63,17 @@ checkpoint's architecture:
        "num_speculative_tokens": 6
      }'
 
-The trained draft config must declare ``DFlash2DraftModel``. Its
-``dflash_config`` must include ``target_layer_ids``, ``conv_group_size``,
-``conv_kernel_size``, ``selector_rank``, and ``selector_top_k``. Set
+NeMo Automodel training exports ``Qwen3DFlash2DraftModel`` so the checkpoint
+can still be reopened by the training stack. The SpeechLM plugin normalizes
+that architecture to vLLM's canonical ``DFlash2DraftModel`` before vLLM wraps
+the speculative config. This is required for vLLM to force its V2 model runner
+and execute the DFlash2 candidate-selector speculator instead of silently
+falling back to plain DFlash. Native configs that already declare
+``DFlash2DraftModel`` remain supported.
+
+The draft's ``dflash_config`` must include ``target_layer_ids``,
+``conv_group_size``, ``conv_kernel_size``, ``selector_rank``, and
+``selector_top_k``. Set
 ``num_speculative_tokens`` to one less than the convolution block size used to
 train the draft: vLLM constructs each runtime block from one anchor plus the
 configured number of draft tokens and does not reject a training/inference
diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py
index add8c2916f31..c8365dd9f011 100644
--- a/nemo/collections/speechlm2/vllm/salm/__init__.py
+++ b/nemo/collections/speechlm2/vllm/salm/__init__.py
@@ -25,10 +25,32 @@
 
 _PKG = "nemo.collections.speechlm2.vllm.salm"
 _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = None
+_AUTOMODEL_DFLASH2_ARCHITECTURES = frozenset(
+    {
+        "Qwen3DFlash2DraftModel",
+        "DFlashQwen3DFlash2DraftModel",
+    }
+)
+
+
+def _normalize_dflash2_architecture(hf_config):
+    """Route Automodel DFlash2 exports to vLLM's canonical runtime.
+
+    Automodel keeps its training class in ``config.json`` so the checkpoint can
+    be reopened for training. The pinned vLLM DFlash2 implementation dispatches
+    both its V2 runner and candidate-selector speculator only when it sees the
+    canonical ``DFlash2DraftModel`` architecture. Normalize before vLLM wraps
+    the draft in ``EAGLEConfig``; otherwise ``method=dflash`` prefixes the
+    Automodel name and silently selects the plain-DFlash speculator.
+    """
+    architectures = getattr(hf_config, "architectures", None) or []
+    if len(architectures) == 1 and architectures[0] in _AUTOMODEL_DFLASH2_ARCHITECTURES:
+        hf_config.architectures = ["DFlash2DraftModel"]
+    return hf_config
 
 
 def _nemo_speechlm_mtp_hf_config_override(hf_config):
-    """Apply the SpeechLM MTP rewrite, then defer unrelated configs to vLLM.
+    """Apply SpeechLM speculative-config rewrites, then defer to vLLM.
 
     This function must remain at module scope: vLLM retains it on the draft
     ``ModelConfig``, which can cross a spawned process boundary. The original
@@ -88,16 +110,17 @@ def _nemo_speechlm_mtp_hf_config_override(hf_config):
         if current_override is _nemo_speechlm_mtp_hf_config_override:
             raise RuntimeError("NeMo SpeechLM MTP override was installed without preserving vLLM's original hook.")
         _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
-    return _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE(hf_config)
+    hf_config = _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE(hf_config)
+    return _normalize_dflash2_architecture(hf_config)
 
 
 _nemo_speechlm_mtp_hf_config_override._nemo_speechlm_mtp_override = True
 
 
 def _patch_vllm_for_nemo_speechlm_mtp() -> None:
-    """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
+    """Extend vLLM's speculative-decoding framework for SpeechLM drafts.
 
-    Three patches are applied on the supported vLLM 0.19+ releases:
+    Four patches are applied on supported vLLM releases:
 
     1. ``MTPModelTypes`` — the Literal type that guards the MTP detection
        branch in ``SpeculativeConfig.__post_init__`` is extended to include
@@ -111,6 +134,10 @@ def _patch_vllm_for_nemo_speechlm_mtp() -> None:
 
     3. ``ModelRegistry`` — ``NeMoSpeechLMMTPModel`` is registered so that
        vLLM can resolve and instantiate it as the draft model.
+
+    4. Automodel DFlash2 architecture names are normalized to vLLM's canonical
+       ``DFlash2DraftModel`` before ``EAGLEConfig`` wrapping, which activates
+       the V2 model runner and candidate-selector speculator.
     """
     from typing import Literal, get_args
 
@@ -166,5 +193,11 @@ def register():
         for automodel_arch in ("Qwen3DFlashDraftModel", "DFlashQwen3DFlashDraftModel"):
             if automodel_arch not in supported_archs:
                 ModelRegistry.register_model(automodel_arch, native_dflash_model_ref)
+    if "DFlash2DraftModel" in supported_archs:
+        native_dflash2_model = ModelRegistry.models["DFlash2DraftModel"]
+        native_dflash2_model_ref = f"{native_dflash2_model.module_name}:{native_dflash2_model.class_name}"
+        for automodel_arch in _AUTOMODEL_DFLASH2_ARCHITECTURES:
+            if automodel_arch not in supported_archs:
+                ModelRegistry.register_model(automodel_arch, native_dflash2_model_ref)
 
     _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 036bf476aeac..b17a5e8d1540 100644
--- a/nemo/collections/speechlm2/vllm/salm/audio.py
+++ b/nemo/collections/speechlm2/vllm/salm/audio.py
@@ -112,9 +112,7 @@ def _build_pe_encoder_from_config(pe_encoder_config: Mapping[str, object]) -> nn
     """Reconstruct a PE encoder whose config and weights are embedded in an HF export."""
     from omegaconf import OmegaConf
 
-    from nemo.collections.asr.modules.parallel_expert_encoder import (
-        ParallelExpertEncoder,
-    )
+    from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoder
 
     if hasattr(pe_encoder_config, "to_dict"):
         pe_encoder_config = pe_encoder_config.to_dict()
diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py
index 4a7bc2760e20..8e228f3f5f86 100644
--- a/tests/collections/speechlm2/test_vllm_plugin.py
+++ b/tests/collections/speechlm2/test_vllm_plugin.py
@@ -55,6 +55,34 @@ def _identity_hf_config_override(hf_config):
     return hf_config
 
 
+@pytest.mark.skipif(not _HAS_CONFIG, reason="SpeechLM vLLM plugin not available")
+@pytest.mark.parametrize(
+    "architecture",
+    ("Qwen3DFlash2DraftModel", "DFlashQwen3DFlash2DraftModel"),
+)
+def test_normalizes_automodel_dflash2_architecture(architecture):
+    """Normalization does not need vLLM installed or model weights loaded."""
+    config = SimpleNamespace(architectures=[architecture])
+
+    result = _salm_module._normalize_dflash2_architecture(config)
+
+    assert result is config
+    assert config.architectures == ["DFlash2DraftModel"]
+
+
+@pytest.mark.skipif(not _HAS_CONFIG, reason="SpeechLM vLLM plugin not available")
+@pytest.mark.parametrize(
+    "architectures",
+    (None, [], ["DFlash2DraftModel"], ["Qwen3DFlashDraftModel"], ["Qwen3DFlash2DraftModel", "Other"]),
+)
+def test_dflash2_normalization_preserves_unrelated_architectures(architectures):
+    config = SimpleNamespace(architectures=architectures)
+
+    _salm_module._normalize_dflash2_architecture(config)
+
+    assert config.architectures == architectures
+
+
 @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
 class TestNeMoSpeechLMConfig:
     """Tests for NeMoSpeechLMConfig."""
@@ -827,6 +855,41 @@ def test_registers_automodel_dflash_architecture_alias(self, monkeypatch):
                 native_model.class_name,
             )
 
+    @pytest.mark.parametrize(
+        "automodel_arch",
+        ("Qwen3DFlash2DraftModel", "DFlashQwen3DFlash2DraftModel"),
+    )
+    def test_routes_automodel_dflash2_to_candidate_selector_runtime(self, monkeypatch, automodel_arch):
+        """Automodel exports must retain DFlash2 semantics after vLLM config wrapping."""
+        from transformers import AutoConfig
+        from vllm.model_executor.models.registry import ModelRegistry
+        from vllm.transformers_utils.configs.eagle import EAGLEConfig
+
+        if "DFlash2DraftModel" not in ModelRegistry.get_supported_archs():
+            pytest.skip("installed vLLM does not provide the DFlash2 candidate-selector runtime")
+
+        monkeypatch.setattr(
+            AutoConfig,
+            "from_pretrained",
+            lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()),
+        )
+
+        register()
+
+        native_model = ModelRegistry.models["DFlash2DraftModel"]
+        automodel_config = AutoConfig.for_model("qwen3", architectures=[automodel_arch])
+        normalized_config = SpeculativeConfig.hf_config_override(automodel_config)
+        runtime_config = EAGLEConfig(normalized_config, method="dflash", model_type="eagle")
+
+        assert normalized_config.architectures == ["DFlash2DraftModel"]
+        assert runtime_config.architectures == ["DFlash2DraftModel"]
+        for alias in (automodel_arch, *runtime_config.architectures):
+            alias_model = ModelRegistry.models[alias]
+            assert (alias_model.module_name, alias_model.class_name) == (
+                native_model.module_name,
+                native_model.class_name,
+            )
+
     def test_model_advertises_eagle3_support(self):
         from vllm.model_executor.models.interfaces import supports_eagle3