Add DFlash speculative decoding support to NeMo SpeechLM vLLM plugin

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

docs/source/speechlm2/vllm_dflash.rst

@@ -0,0 +1,95 @@
+DFlash speculative decoding with vLLM
+======================================
+
+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
+------------
+
+* 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.
+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
+-----------------
+
+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 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@3406ec1dae9916f920b90f0dbf90dcf54923d042"
+
+   vllm serve /path/to/vllm-ready-speechlm-checkpoint \
+     --trust-remote-code \
+     --speculative-config '{
+       "method": "dflash",
+       "model": "/path/to/trained-lightning-dflash2-checkpoint",
+       "num_speculative_tokens": 6
+     }'
+
+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
+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
+----------
+
+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.

examples/speechlm2/to_hf.py

@@ -102,6 +102,37 @@ 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)
+    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
+        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
@@ -116,9 +147,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)
@@ -135,16 +165,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:
@@ -160,8 +192,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:
@@ -175,10 +210,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)
@@ -195,15 +232,27 @@ 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)
-    arch = _detect_vllm_architecture(arch_model_cfg)
+        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
-    config_path.write_text(json.dumps(config, indent=2) + "\n")
+    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)
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [
@@ -213,9 +262,20 @@ 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, 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)
+    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.

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,