feat(parakeet): add Parakeet TDT support

benchmarks/performance/baselines/task_reference.py

@@ -562,7 +562,7 @@ def _disable_nemo_asr_cuda_graphs(model: Any) -> bool:
 def _load_asr(
     arguments: argparse.Namespace,
     request: Mapping[str, Any],
-    _options: Mapping[str, Any],
+    options: Mapping[str, Any],
 ) -> Session:
     import torch
     from tools.validation.engine import _read_wav_float32, _resample_audio
@@ -615,11 +615,25 @@ def invoke() -> Mapping[str, Any]:
             return {"text": _transcription_text(result), "output_tokens": None}
 
     else:
-        from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
+        import transformers
+        from transformers import AutoProcessor
+
+        auto_model_class = str(
+            options.get("auto_model_class", "AutoModelForSpeechSeq2Seq")
+        )
+        if auto_model_class not in {
+            "AutoModelForSpeechSeq2Seq",
+            "AutoModelForTDT",
+        }:
+            raise ValueError(
+                "hf-transformers-asr adapter_options.auto_model_class must be "
+                "AutoModelForSpeechSeq2Seq or AutoModelForTDT"
+            )
+        model_loader = getattr(transformers, auto_model_class)
 
         processor = AutoProcessor.from_pretrained(arguments.model, **_processor_kwargs(arguments))
         model = (
-            AutoModelForSpeechSeq2Seq.from_pretrained(
+            model_loader.from_pretrained(
                 arguments.model, **_load_kwargs(arguments, torch)
             )
             .eval()
@@ -638,11 +652,19 @@ def invoke() -> Mapping[str, Any]:
 
         def invoke() -> Mapping[str, Any]:
             with torch.inference_mode():
-                generated = model.generate(**inputs, max_new_tokens=max_new_tokens)
+                generate_options: dict[str, Any] = {"max_new_tokens": max_new_tokens}
+                if auto_model_class == "AutoModelForTDT":
+                    generate_options["return_dict_in_generate"] = True
+                generated = model.generate(**inputs, **generate_options)
             sequences = generated.sequences if hasattr(generated, "sequences") else generated
             token_ids = [int(token) for token in sequences[0].detach().cpu().tolist()]
+            if auto_model_class == "AutoModelForTDT":
+                decoded = processor.decode(sequences, skip_special_tokens=True)
+            else:
+                decoded = processor.batch_decode(sequences, skip_special_tokens=True)
+            text = decoded[0] if isinstance(decoded, (list, tuple)) else str(decoded)
             return {
-                "text": processor.batch_decode(sequences, skip_special_tokens=True)[0],
+                "text": text,
                 "token_ids": token_ids,
                 "output_tokens": len(token_ids),
             }

benchmarks/performance/baselines/timing_contracts.py

@@ -16,6 +16,7 @@
     "elf_flow",
     "internvl",
     "locateanything",
+    "parakeet_tdt",
     "patchtsmixer",
     "patchtst",
     "phi4_multimodal",

benchmarks/performance/release.yaml

@@ -631,6 +631,21 @@ entries:
       runner: hf-transformers
       mode: torch-compile
       compile_scope: model.forward
+  - id: parakeet_tdt.transcribe
+    family: parakeet_tdt
+    operation: transcribe
+    model: parakeet-tdt-0.6b-v3
+    workload:
+      testcase: parakeet-tdt-0.6b-v3
+    baseline:
+      runner: task-reference
+      adapter: hf-transformers-asr
+      mode: hf-eager
+      reference_backend: hf_transformers
+      timing_scope: task-model-call-wall
+      input_preparation_included: false
+      adapter_options:
+        auto_model_class: AutoModelForTDT
   - id: opt.generate
     family: opt
     operation: generate

python/tensorrt_model_connect/families/__init__.py

@@ -45,6 +45,7 @@ class _FamilyMetadata:
     diffusion_pipeline_classes: frozenset[str]
     nemo_target_patterns: frozenset[str]
     nemo_model_type: str
+    model_resolution_priority: int = 0
     nemo_archive_adapter: str = ""
     hf_allow_patterns: tuple[str, ...] = ()
     hf_required_files: tuple[str, ...] = ()
@@ -169,6 +170,10 @@ def _load_family_metadata() -> list[_FamilyMetadata]:
             ),
             nemo_model_type=raw.get("nemo_model_type", "")
             if isinstance(raw.get("nemo_model_type"), str) else "",
+            model_resolution_priority=raw.get("model_resolution_priority", 0)
+            if isinstance(raw.get("model_resolution_priority", 0), int)
+            and not isinstance(raw.get("model_resolution_priority", 0), bool)
+            else 0,
             nemo_archive_adapter=raw.get("nemo_archive_adapter", "")
             if isinstance(raw.get("nemo_archive_adapter"), str) else "",
             hf_allow_patterns=tuple(_metadata_strings(raw.get("hf_allow_patterns"))),
@@ -201,6 +206,14 @@ def _load_family_metadata() -> list[_FamilyMetadata]:
     return metadata
 
 
+def _model_resolution_metadata() -> list[_FamilyMetadata]:
+    """Return source adapters from most specific to broadest family claim."""
+    return sorted(
+        _load_family_metadata(),
+        key=lambda candidate: (-candidate.model_resolution_priority, candidate.id),
+    )
+
+
 def _add_index_value(
     index: dict[str, list[_ModuleCandidate]],
     key: str,
@@ -560,7 +573,7 @@ def resolve_config_from_model_dir(model_dir: str | Path) -> dict[str, Any] | Non
 def resolve_family_model_dir(model_dir: str | Path) -> str | None:
     """Ask family adapters to stage a non-flat model repository."""
     path = Path(model_dir)
-    for meta in _load_family_metadata():
+    for meta in _model_resolution_metadata():
         if not meta.model_dir_adapter:
             continue
         adapter = _load_metadata_callable_from_file(
@@ -600,7 +613,7 @@ def family_prefers_native_default_build(
 def resolve_nemo_archive_model_dir(nemo_path: str | Path) -> str | None:
     """Ask family-owned NeMo archive adapters to synthesize a model dir."""
     path = Path(nemo_path)
-    for meta in _load_family_metadata():
+    for meta in _model_resolution_metadata():
         if not meta.nemo_archive_adapter:
             continue
         adapter = _load_metadata_callable_from_file(meta, meta.nemo_archive_adapter)
@@ -719,7 +732,7 @@ def resolve_nemo_model_type(config: dict) -> str:
     """
     target = str(config.get("target", "") or config.get("_target_", ""))
     target_key = target.lower()
-    for meta in _load_family_metadata():
+    for meta in _model_resolution_metadata():
         if not meta.nemo_model_type or not meta.nemo_target_patterns:
             continue
         for pattern in meta.nemo_target_patterns:

python/tensorrt_model_connect/families/parakeet_tdt/MODEL.toml

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+id = "parakeet_tdt"
+plugin = "parakeet_tdt"
+python_profile_specs = [
+  "parakeet_tdt_reference|families/parakeet_tdt/python_profile_requirements/parakeet_tdt_reference.lock.txt|families/parakeet_tdt/python_profile_verify.py|true",
+]
+default_execution_profiles = [
+  "reference|parakeet_tdt_reference",
+]
+module = "plugin"
+model_dir_adapter = "nemo_archive.py|resolve_model_dir"
+nemo_archive_adapter = "nemo_archive.py|resolve_nemo_archive"
+model_resolution_priority = 100
+nemo_model_type = "parakeet_tdt"
+nemo_target_patterns = ["EncDecRNNTBPEModel", "TDT", "tdt"]
+aliases = ["parakeet_tdt", "parakeet-tdt", "parakeet_tdt_0_6b_v3"]
+prefixes = ["parakeet_tdt", "parakeet-tdt"]

python/tensorrt_model_connect/families/parakeet_tdt/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .plugin import ParakeetTDTPlugin, plugin
+
+__all__ = ["ParakeetTDTPlugin", "plugin"]

python/tensorrt_model_connect/families/parakeet_tdt/checkpoint.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Checkpoint readers and tensor normalization for Parakeet TDT."""
+
+from __future__ import annotations
+
+import io
+import tarfile
+from pathlib import Path
+from typing import Mapping
+
+import numpy as np
+
+
+class WeightDict(dict):
+    """Normalized family-owned build tensors."""
+
+
+def _array(value) -> np.ndarray:
+    if hasattr(value, "detach"):
+        value = value.detach().cpu().numpy()
+    return np.ascontiguousarray(np.asarray(value, dtype=np.float32))
+
+
+def _transpose(value) -> np.ndarray:
+    array = _array(value)
+    if array.ndim != 2:
+        raise ValueError(f"expected rank-2 weight, got {array.shape}")
+    return np.ascontiguousarray(array.T)
+
+
+def _transpose_2d(value, name: str, precision: str = "fp32") -> np.ndarray:
+    """Family graph-builder adapter with the repository's mapper signature."""
+    del name, precision
+    return _transpose(value)
+
+
+def map_transducer_weights(
+    state: Mapping[str, object],
+    *,
+    vocab_size: int,
+    duration_count: int,
+    decoder_layers: int,
+    decoder_hidden_size: int,
+    encoder_hidden_size: int,
+) -> WeightDict:
+    """Normalize the predictor/projector/joint tensors shared by HF and NeMo layouts."""
+    out = WeightDict()
+    out["pred_embedding"] = _array(state["decoder.embedding.weight"])
+    for layer in range(decoder_layers):
+        prefix = "decoder.lstm"
+        w_ih = _array(state[f"{prefix}.weight_ih_l{layer}"])
+        w_hh = _array(state[f"{prefix}.weight_hh_l{layer}"])
+        b_ih = _array(state[f"{prefix}.bias_ih_l{layer}"])
+        b_hh = _array(state[f"{prefix}.bias_hh_l{layer}"])
+        expected = (4 * decoder_hidden_size, decoder_hidden_size)
+        if w_ih.shape != expected or w_hh.shape != expected:
+            raise ValueError(
+                f"predictor layer {layer} has shapes {w_ih.shape}/{w_hh.shape}, expected {expected}"
+            )
+        out[f"pred.{layer}.w_ih_t"] = np.ascontiguousarray(w_ih.T)
+        out[f"pred.{layer}.w_hh_t"] = np.ascontiguousarray(w_hh.T)
+        out[f"pred.{layer}.bias"] = np.ascontiguousarray(b_ih + b_hh)
+
+    out["decoder_projector_w"] = _transpose(state["decoder.decoder_projector.weight"])
+    out["decoder_projector_b"] = _array(state["decoder.decoder_projector.bias"])
+    out["encoder_projector_w"] = _transpose(state["encoder_projector.weight"])
+    out["encoder_projector_b"] = _array(state["encoder_projector.bias"])
+
+    joint_w = _array(state["joint.head.weight"])
+    joint_b = _array(state["joint.head.bias"])
+    expected_outputs = vocab_size + duration_count
+    if joint_w.shape != (expected_outputs, decoder_hidden_size):
+        raise ValueError(
+            f"joint.head.weight has shape {joint_w.shape}, expected "
+            f"{(expected_outputs, decoder_hidden_size)}"
+        )
+    if joint_b.shape != (expected_outputs,):
+        raise ValueError(
+            f"joint.head.bias has shape {joint_b.shape}, expected {(expected_outputs,)}"
+        )
+    out["joint_token_w"] = np.ascontiguousarray(joint_w[:vocab_size].T)
+    out["joint_token_b"] = np.ascontiguousarray(joint_b[:vocab_size])
+    out["joint_duration_w"] = np.ascontiguousarray(joint_w[vocab_size:].T)
+    out["joint_duration_b"] = np.ascontiguousarray(joint_b[vocab_size:])
+    out["_encoder_hidden"] = encoder_hidden_size
+    out["_pred_hidden"] = decoder_hidden_size
+    out["_pred_layers"] = decoder_layers
+    out["_vocab"] = vocab_size
+    out["_duration_count"] = duration_count
+    return out
+
+
+def load_hf_safetensors(model_dir: str | Path) -> dict[str, np.ndarray]:
+    from safetensors.numpy import load_file
+
+    path = Path(model_dir) / "model.safetensors"
+    if not path.is_file():
+        raise FileNotFoundError(f"Parakeet HF checkpoint is missing {path.name}")
+    return dict(load_file(path))
+
+
+def load_nemo_archive(model_dir: str | Path) -> tuple[dict, dict]:
+    import torch
+    import yaml
+
+    root = Path(model_dir)
+    archive = root if root.suffix == ".nemo" else next(iter(sorted(root.glob("*.nemo"))), None)
+    if archive is None:
+        raise FileNotFoundError(f"No .nemo file found in {root}")
+    state = config = None
+    with tarfile.open(archive, "r") as tar:
+        for member in tar.getmembers():
+            name = Path(member.name).name
+            if name == "model_config.yaml":
+                extracted = tar.extractfile(member)
+                if extracted is not None:
+                    config = yaml.safe_load(extracted.read())
+            elif name == "model_weights.ckpt":
+                extracted = tar.extractfile(member)
+                if extracted is not None:
+                    state = torch.load(
+                        io.BytesIO(extracted.read()), map_location="cpu", weights_only=True
+                    )
+    if not isinstance(config, dict):
+        raise FileNotFoundError(f"model_config.yaml not found in {archive}")
+    if not isinstance(state, dict):
+        raise FileNotFoundError(f"model_weights.ckpt not found in {archive}")
+    if isinstance(state.get("state_dict"), dict):
+        state = state["state_dict"]
+    return state, config