[Alex-Wengg]

Overview

Adds Nemotron3DiarizationForAudioFrameClassification for nvidia/Nemotron-3-Diarization, a 99M-parameter streaming Sortformer speaker diarizer. For every 10 ms frame of a recording it returns the logits of up to 8 speakers. vLLM serves it offline: one whole recording per request, through the plugin pooling task. The model is tokenizer-free and raw-input-only, like Terratorch, and wraps the Transformers module as is (Transformers ≥ 5.18).

No existing issue or PR: gh pr list / gh issue list searches for Nemotron-3-Diarization, Nemotron3Diarization, diarization and sortformer return only unrelated ASR/realtime items.

Claims

• LLM.encode(..., pooling_task="plugin") returns [num_frames, 8] logits per recording. With fp32 these are bit-identical to Transformers Nemotron3DiarizationForAudioFrameClassification given the same features, including recordings long enough to run the speaker cache over 6 chunks.
• Recordings of different lengths in one encode call give the same output as single requests.
• A prompt without audio is rejected with a VLLMValidationError instead of returning empty logits or crashing the engine.
• The engine starts with default settings: vLLM selects eager mode and the V1 model runner for this model on its own.

Validation

NVIDIA RTX A6000, Transformers 5.22.0 (the CI pin), torch 2.14.0, fp32. Clips: the pyannote 2-speaker sample.wav (30 s), its first 17.3 s, and five copies back to back (150 s, 6 offline chunks).

clip	frames	max |Δlogit| vs Transformers	extract_speaker_dict segments
30 s	3001	0	identical
150 s	15001	0	identical
17.3 s	1731	0	identical

• Batching: [150 s, 17.3 s, 30 s, 17.3 s] in one call: every output equals its single-request result exactly (the batch runs in ~0.9 s).
• Against the reference RTTM (30 s clip): speaker onsets 6.72 s / 7.56 s vs 6.69 s / 7.55 s; frame-level DER without collar 2.38% (miss 0.57%, false alarm 1.81%, confusion 0%).
• Half precision: dtype="auto" resolves to float16 for this fp32 checkpoint. fp16 and bf16 keep 99.98–100% frame agreement and the same onsets, but move a few segment boundaries, so the docs recommend dtype="float32".
• The example examples/pooling/plugin/nemotron3_diarization_offline.py --audio sample.wav prints 11 segments against the RTTM's 10. Boundaries are within 0.05 s of the RTTM except for speaker 1 at two points: at 14.10 s an extra 0.06 s segment plus an onset 0.20 s early (14.29 s vs 14.49 s), and at 18.07 s an onset 0.08 s early.

Exact parity holds when both sides use identical features. vLLM's frontend computes features with one torch thread, and this feature extractor's CPU output varies by ~1e-6 between 1 and 6 threads. Against a default-threads Transformers run the max |Δlogit| is 2.5e-3, with identical segments. The reference above therefore uses torch.set_num_threads(1).

Reproduce

import soundfile as sf, torch
from transformers import AutoModelForAudioFrameClassification, AutoProcessor
from vllm import LLM

model_id = "nvidia/Nemotron-3-Diarization"
wav, sr = sf.read("sample.wav", dtype="float32")

# Transformers, features computed with one thread as in vLLM's frontend
torch.set_num_threads(1)
processor = AutoProcessor.from_pretrained(model_id)
hf = AutoModelForAudioFrameClassification.from_pretrained(model_id, dtype=torch.float32).cuda()
with torch.inference_mode():
    ref = hf(**processor(wav, sampling_rate=sr).to("cuda")).logits[0].cpu()

# vLLM
llm = LLM(model=model_id, runner="pooling", skip_tokenizer_init=True, dtype="float32")
(out,) = llm.encode(
    [{"prompt_token_ids": [0], "multi_modal_data": {"audio": (wav, sr)}}],
    pooling_task="plugin",
)
print((out.outputs.data - ref).abs().max())  # tensor(0.)

Tests

pytest tests/models/test_registry.py -k Nemotron3Diarization        # 1 passed
pytest tests/models/test_initialization.py -k Nemotron3Diarization  # 1 passed (locally with gpu_memory_utilization=0.15, a shared GPU)
pytest tests/test_config.py                                         # 496 passed, 7 skipped, 3 failed*
pytest tests/entrypoints/llm/test_prompt_validation.py              # 5 passed (incl. the new test)
pytest tests/model_executor/test_routed_experts_capture.py          # 51 passed
pre-commit run --files <changed files>                              # passed

* test_rope_customization, test_is_encoder_decoder[meta-llama/Llama-3.2-1B-Instruct-False] and test_eagle_draft_model_config fail with "You are trying to access a gated repo": they need gated meta-llama checkpoints, and the box has no HF token. The new test_v2_model_runner_raw_input_only_model_falls_back_to_v1 passes. No model-specific correctness test is added; the registry entry covers loading.

Details

• Model file follows terratorch.py. The processor overrides apply() to run the checkpoint's NemotronAsrStreamingFeatureExtractor, as Nemotron3DiarizationProcessor does offline. It returns the placeholder prompt [0] with a zero-length audio placeholder. The model loops over the step's recordings and runs the Transformers offline forward on each, under gpu_sync_allowed(). That forward chunks the recording and updates the speaker cache, which reads back to the host.
• Padded output: V1's _pool needs a tensor (it reads hidden_states.device), so the step's logits are padded to [num_reqs, max_frames, 8]. A small IdentityPooler subclass trims each one back to its own frame count.
• RoPE frequencies are recomputed on the CPU after construction. vLLM builds the module on the GPU, where inv_freq comes out ~4e-9 different. That alone shifts the logits by up to 2.4e-3 on the 150 s clip, as the speaker cache carries small differences from one chunk to the next.
• Placeholder token: the arch-config converter declares a one-token vocabulary. Since [Bugfix][Core] Validate prompt token ids when skip_tokenizer_init is set #62911, prompt ids are validated against the model vocab even with skip_tokenizer_init, and this config has no vocab_size.
• Eager mode: a verify_and_update_model_config hook sets enforce_eager=True. The per-chunk host syncs break CUDA graph capture (device-side assert).
• Model Runner V2: MRV2 has no raw-multimodal-input path and no plugin pooling task. So _get_v2_model_runner_unsupported_features() now lists raw-input-only models, and they fall back to V1 with the usual warning instead of failing at load. Terratorch, the other raw-input-only model, no longer needs VLLM_USE_V2_MODEL_RUNNER=0 either, though it is separately affected by [Bugfix][Core] Validate prompt token ids when skip_tokenizer_init is set #62911. Supporting these models on MRV2 is a separate change.
• Prompts without audio: a token prompt without multi_modal_data skips the multimodal processor, so InputProcessor now rejects it for raw-input-only models. Before, such a prompt returned empty logits on its own, and batched with real audio it crashed the engine (pooler output length mismatch). The new test in test_prompt_validation.py covers it.
• Memory: a step runs up to max_num_seqs recordings with all their features on the GPU (about 200 MB per hour of audio in fp32). The docs suggest lowering max_num_seqs for long recordings.
• Registry: listed with the embedding models next to Terratorch, the other raw-input plugin-task model. The test entry uses require_embed_inputs=True to skip tokenizer init, as Terratorch does.
• Follow-ups (out of scope): online serving (/pooling cannot carry audio for a tokenizer-free model without an IO processor plugin); streaming sessions that keep the speaker cache across requests; MRV2 support.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--65002.org.readthedocs.build/en/65002/

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[DarkLight1337]

This is quite a small model. I wonder how much performance improvement does vLLM actually provide here over other baselines?

[Alex-Wengg]

@DarkLight1337 good point, honestly not much.

should I move it to an out-of-tree plugin, like vllm-parakeet-plugin (#65005) or just close this PR . support nemotron 3 on vllm in a form for devs that want easier time customizing different models easily

[DarkLight1337]

If you are ok with maintaining it and feel that it's valuable then you could do plugin. Either way let's close this PR.