feat(speechlm2): add native MTP speculative decoding in vLLM

examples/speechlm2/to_hf.py

@@ -27,6 +27,7 @@
 from safetensors.torch import save_file
 
 from nemo.collections.speechlm2.parts.hf_hub import LLM_BACKBONE_DIR
+from nemo.collections.speechlm2.vllm.salm.config import _mtp_pattern_from_backbone_config, _resolve_speechlm_mtp_config
 from nemo.core.classes.common import safe_instantiate
 from nemo.core.config import hydra_runner
 from nemo.utils.dtype import str_to_dtype
@@ -186,6 +187,77 @@ def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[
     config["dtype"] = dtype_name
     config["torch_dtype"] = dtype_name
 
+    llm = getattr(model, "llm", None)
+    text_config = getattr(llm, "config", None)
+    explicit_mtp = config.get("mtp")
+    mtp_enabled = (
+        bool(explicit_mtp.get("enabled", True))
+        if isinstance(explicit_mtp, dict)
+        else bool(config.get("compute_mtp", False))
+    )
+    runtime_mtp_config = getattr(llm, "mtp_config", None)
+    runtime_mtp_depth = getattr(runtime_mtp_config, "num_layers", None)
+    if runtime_mtp_depth is None and text_config is not None:
+        runtime_mtp_depth = getattr(text_config, "num_nextn_predict_layers", 0)
+    if not isinstance(explicit_mtp, dict) and mtp_enabled and not int(runtime_mtp_depth or 0):
+        # compute_mtp is the legacy switch. Modern SALMAutomodel recipes
+        # suppress a checkpoint-native head when no explicit mtp block is
+        # present; do not let a stale flag recreate a serving-only MTP module.
+        config["compute_mtp"] = False
+        mtp_enabled = False
+
+    if text_config is not None and mtp_enabled:
+        use_repeated_layer = getattr(
+            runtime_mtp_config,
+            "use_repeated_layer",
+            bool(explicit_mtp.get("use_repeated_layer", False)) if isinstance(explicit_mtp, dict) else False,
+        )
+        resolved_mtp = dict(explicit_mtp) if isinstance(explicit_mtp, dict) else {}
+
+        raw_mtp_pattern = getattr(text_config, "mtp_hybrid_override_pattern", None)
+        mtp_block_types = getattr(text_config, "mtp_layers_block_type", None)
+        if raw_mtp_pattern is not None:
+            if not isinstance(raw_mtp_pattern, str) or (not raw_mtp_pattern and not mtp_block_types):
+                raise ValueError(
+                    f"Built LLM has invalid mtp_hybrid_override_pattern={raw_mtp_pattern!r}; " "cannot export it."
+                )
+        actual_mtp_pattern = _mtp_pattern_from_backbone_config(text_config)
+        if actual_mtp_pattern is not None:
+            # 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 serving constructs matching
+            # physical layers.
+            resolved_mtp["hybrid_override_pattern"] = actual_mtp_pattern
+
+        actual_mtp_depth = getattr(text_config, "num_nextn_predict_layers", None)
+        logical_mtp_depth = getattr(runtime_mtp_config, "num_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 use_repeated_layer:
+                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.
+                logical_mtp_depth = actual_mtp_depth
+
+        config["mtp"] = _resolve_speechlm_mtp_config(
+            mtp=resolved_mtp,
+            compute_mtp=bool(config.get("compute_mtp", False)),
+            text_config=text_config,
+            num_nextn_predict_layers=logical_mtp_depth,
+            use_repeated_layer=use_repeated_layer,
+        )
+    elif mtp_enabled and isinstance(explicit_mtp, dict):
+        raise ValueError(
+            "The root mtp config enables MTP, but the instantiated model has no positive-depth MTP head to export."
+        )
     return config
 
 
@@ -309,6 +381,7 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     else:
         config.pop("llm_config", None)
     config.pop("audio_token_index", None)
+    config.pop("image_token_index", None)
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [

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

@@ -25,6 +25,113 @@
 """
 
 _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 closure 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: 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 = int(mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0) or 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(
+                {
+                    # vLLM instantiates one physical prediction step and reuses
+                    # it for arbitrary speculative K. Repeated-layer training
+                    # produces exactly that checkpoint layout.
+                    "n_predict": 1,
+                    "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.
+        import vllm.config.speculative as _spec_mod
+
+        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
+    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:
+    """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
+
+    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
+       ``"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 enabled 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
+
+    # 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",)]
+
+    # Route SpeechLM MTP checkpoints through 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.
+        # Replacing it during later registration could capture a third-party
+        # wrapper that already delegates to us and create an override cycle.
+        if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
+            _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_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
+
+    ModelRegistry.register_model(
+        "NeMoSpeechLMMTPModel",
+        f"{_PKG}.mtp:NeMoSpeechLMMTP",
+    )
 
 
 def register():
@@ -52,6 +159,8 @@ def register():
 
     MODELS_CONFIG_MAP["NeMoSpeechLMForConditionalGeneration"] = NeMoSpeechLMForConditionalGenerationConfig
 
+    _patch_vllm_for_nemo_speechlm_mtp()
+
     from nemo.collections.speechlm2.vllm.salm.runtime_compat import install_prompt_contract
 
     install_prompt_contract()