[Model] Add NVIDIA Nemotron-3-Diarization (offline speaker diarization)

docs/models/pooling_models/specific_models.md

@@ -484,3 +484,46 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
 `embed&token_classify` task is an internal execution contract for this plugin,
 not a generic Pooling API response format. Without the plugin, select one of
 the three concrete tasks above.
+
+## NVIDIA Nemotron-3-Diarization
+
+`nvidia/Nemotron-3-Diarization` is a streaming Sortformer speaker diarization
+model. For every 10 ms frame of a recording it predicts the activity of up to 8
+speakers, numbered in order of first arrival.
+
+| Architecture | Inputs | Example HF Models |
+| - | - | - |
+| `Nemotron3DiarizationForAudioFrameClassification` | A | `nvidia/Nemotron-3-Diarization` |
+
+vLLM runs it offline, one whole recording per request, through the `plugin`
+pooling task. The model takes no text, so start the engine with
+`skip_tokenizer_init=True`. vLLM runs it in eager mode on the V1 model runner,
+and its feature extractor needs the `librosa` package. Use `dtype="float32"`,
+the checkpoint's dtype, for logits that match Transformers; half precision moves
+a few segment boundaries. Each output is the `[num_frames, 8]` logits; turn
+them into speech segments with the Transformers processor:
+
+```python
+from transformers import AutoProcessor
+from vllm import LLM
+
+model = "nvidia/Nemotron-3-Diarization"
+llm = LLM(model=model, runner="pooling", skip_tokenizer_init=True, dtype="float32")
+processor = AutoProcessor.from_pretrained(model)
+
+(output,) = llm.encode(
+    [{"prompt_token_ids": [0], "multi_modal_data": {"audio": (audio, sample_rate)}}],
+    pooling_task="plugin",
+)
+segments = processor.extract_speaker_dict(output.outputs.data[None])[0]
+# [{'Start': <seconds>, 'End': <seconds>, 'Speaker': <index>}, ...]
+```
+
+A step runs up to `max_num_seqs` recordings, and the features of all of them sit
+on the GPU together: about 200 MB per hour of audio in fp32. For long recordings,
+lower `max_num_seqs` to bound this memory.
+
+Streaming sessions, which carry the speaker cache from one request to the next,
+and the online Pooling API are not supported yet. See
+[examples/pooling/plugin/nemotron3_diarization_offline.py](../../../examples/pooling/plugin/nemotron3_diarization_offline.py)
+for a complete example.

examples/pooling/plugin/nemotron3_diarization_offline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline speaker diarization with NVIDIA Nemotron-3-Diarization.
+
+Each request is one whole recording. The model returns the logits of up to 8
+speakers for every 10 ms frame, `[num_frames, 8]`, and the Transformers processor
+turns them into speech segments.
+
+Usage::
+
+    python nemotron3_diarization_offline.py --audio meeting.wav
+"""
+
+from argparse import Namespace
+
+from transformers import AutoProcessor
+
+from vllm import LLM, EngineArgs
+from vllm.assets.audio import AudioAsset
+from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+
+def parse_args():
+    parser = FlexibleArgumentParser()
+    parser = EngineArgs.add_cli_args(parser)
+    parser.add_argument(
+        "--audio",
+        type=str,
+        default=None,
+        help="Path to the recording to diarize. Defaults to a vLLM audio asset.",
+    )
+    parser.set_defaults(
+        model="nvidia/Nemotron-3-Diarization",
+        runner="pooling",
+        skip_tokenizer_init=True,
+        dtype="float32",
+    )
+    return parser.parse_args()
+
+
+def main(args: Namespace):
+    audio_path = args.__dict__.pop("audio")
+    if audio_path is None:
+        audio, sample_rate = AudioAsset("winning_call").audio_and_sample_rate
+    else:
+        audio, sample_rate = load_audio(audio_path, sr=None)
+
+    llm = LLM(**vars(args))
+    processor = AutoProcessor.from_pretrained(args.model)
+
+    # The model reads the audio only; the prompt is a placeholder token.
+    outputs = llm.encode(
+        [
+            {
+                "prompt_token_ids": [0],
+                "multi_modal_data": {"audio": (audio, sample_rate)},
+            }
+        ],
+        pooling_task="plugin",
+    )
+
+    for output in outputs:
+        logits = output.outputs.data  # [num_frames, num_speakers]
+        for segment in processor.extract_speaker_dict(logits[None])[0]:
+            print(
+                f"speaker {segment['Speaker']}: "
+                f"{segment['Start']:.2f}s - {segment['End']:.2f}s"
+            )
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

tests/entrypoints/llm/test_prompt_validation.py

@@ -37,6 +37,20 @@ def test_out_of_vocab_token_skip_tokenizer_init(vllm_runner):
         runner.llm.generate({"prompt_token_ids": [999999]})
 
 
+def test_raw_input_only_model_requires_mm_data(vllm_runner):
+    # A prompt without multimodal data would give such a model nothing to run on,
+    # so it must be rejected before it reaches the engine.
+    with (
+        vllm_runner(
+            "nvidia/Nemotron-3-Diarization",
+            runner="pooling",
+            skip_tokenizer_init=True,
+        ) as runner,
+        pytest.raises(VLLMValidationError, match="must include `multi_modal_data`"),
+    ):
+        runner.llm.encode({"prompt_token_ids": [0]}, pooling_task="plugin")
+
+
 def test_require_mm_embeds(vllm_runner):
     with (
         vllm_runner(

tests/model_executor/test_routed_experts_capture.py

@@ -761,6 +761,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch):
             use_mla=False,
             logits_processors=None,
             enable_prompt_embeds=False,
+            is_multimodal_raw_input_only_model=False,
         ),
         aux_output_config=SimpleNamespace(enable_return_routed_experts=True),
         speculative_config=None,

tests/models/registry.py

@@ -664,6 +664,11 @@ def check_available_online(
         "nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True
     ),
     "LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
+    "Nemotron3DiarizationForAudioFrameClassification": _HfExamplesInfo(
+        "nvidia/Nemotron-3-Diarization",
+        min_transformers_version="5.18.0",
+        require_embed_inputs=True,
+    ),
     "Phi3VForCausalLM": _HfExamplesInfo(
         "TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
     ),

tests/test_config.py

@@ -1121,12 +1121,31 @@ def test_v2_model_runner_heterogeneous_vocab_draft_falls_back_to_v1(
 def test_v2_model_runner_supports_custom_logits_processors():
     config = VllmConfig()
     config.model_config = cast(
-        ModelConfig, SimpleNamespace(logits_processors=["a.b:C"])
+        ModelConfig,
+        SimpleNamespace(
+            logits_processors=["a.b:C"], is_multimodal_raw_input_only_model=False
+        ),
     )
 
     assert config._get_v2_model_runner_unsupported_features() == []
 
 
+@pytest.mark.parametrize("raw_input_only", [False, True])
+def test_v2_model_runner_raw_input_only_model_falls_back_to_v1(raw_input_only):
+    """Model Runner V2 does not pass raw multimodal inputs to the model, so
+    raw-input-only models (e.g. Terratorch) must stay on V1."""
+    config = VllmConfig()
+    config.model_config = cast(
+        ModelConfig,
+        SimpleNamespace(is_multimodal_raw_input_only_model=raw_input_only),
+    )
+
+    expected = ["raw-input-only multimodal models"]
+    assert config._get_v2_model_runner_unsupported_features() == (
+        expected if raw_input_only else []
+    )
+
+
 @pytest.mark.parametrize(
     "architecture",
     ["DFlash2DraftModel", "DFlash2BeamDraftModel", "LiLiCorrDraftModel"],

vllm/config/vllm.py

@@ -3129,6 +3129,12 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         if self.parallel_config.use_ubatching:
             unsupported.extend(self._get_dbo_unsupported_features())
 
+        if (
+            self.model_config is not None
+            and self.model_config.is_multimodal_raw_input_only_model
+        ):
+            unsupported.append("raw-input-only multimodal models")
+
         return unsupported
 
     def _get_standalone_cpu_model_runner_unsupported_features(self) -> list[str]:

vllm/model_executor/models/config.py

@@ -930,6 +930,16 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         )
 
 
+class Nemotron3DiarizationForAudioFrameClassificationConfig(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
+        # The forward syncs with the host at every chunk of a recording, so it
+        # cannot be captured in a CUDA graph.
+        if not model_config.enforce_eager:
+            logger.info("Nemotron-3-Diarization runs in eager mode.")
+            model_config.enforce_eager = True
+
+
 class NomicBertModelConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -1325,6 +1335,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
     "Nemotron3_5TranscribeForConditionalGeneration": (
         Nemotron3_5TranscribeForConditionalGenerationConfig
     ),
+    "Nemotron3DiarizationForAudioFrameClassification": (
+        Nemotron3DiarizationForAudioFrameClassificationConfig
+    ),
     "NomicBertModel": NomicBertModelConfig,
     "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
     "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,