Enable Nemotron-3-Diarization on MLX, CUDA, XNNPACK, and Vulkan

[mergennachin]

🚀 The feature, motivation and pitch

Enable NVIDIA Nemotron-3-Diarization in ExecuTorch, with export and native inference examples for MLX, CUDA, XNNPACK, and Vulkan. Support offline recordings and streaming audio so applications can produce speaker timestamps locally for meetings, calls, and speech pipelines.

The model card describes a ~100M-parameter model supporting up to eight speakers, ordered by first arrival. It consumes 16 kHz mono audio, uses 128-bin mel features and a 31-layer Transformer with RoPE, stacks features to an 80 ms encoder frame rate, and uses a Conv1D upsampling head to produce per-speaker activity probabilities at a default 10 ms resolution. Streaming preserves speaker identities through an Arrival-Order Speaker Cache (AOSC) and FIFO context.

Progress

MLX support landed in #23137. The Nemotron example includes export/build/run commands, BF16 by default (with an FP32 option), and a native C++ runner for all four offline/streaming presets.

The exporter captures the Transformers implementation with torch.export. The preprocessor uses XNNPACK with CPU fallback kernels, including FFT; pre_encode and encode use MLX. Streaming scheduling and speaker-cache state remain in the native runner. Final-window handling retains the centered-STFT extra frame and rounds to the subsampling factor of eight, with a two-encoder-frame minimum. The previous pad-to-16 behavior was removed.

Validation on Apple Silicon passed BF16 export, native builds, lint, and 73 local smoke/integration checks: the relocated runner, integration and probability comparisons for all four presets, and 64 tail cases covering all 16 frame-count remainders across those presets. Removing pad-to-16 left the eight-second regression clip's probabilities exactly unchanged in all four presets. The tail tests cover streaming partitioning, reset, output lengths, and short/final inputs; they are not a full DER evaluation. MLX build CI is included; evaluation artifacts remain outside the repository.

CUDA, Vulkan, and XNNPACK neural inference remain open. XNNPACK preprocessing in the MLX example does not complete the CPU backend workstream. Dependency/checkpoint pinning and the broader acceptance criteria below also remain open.

Proposed scope

• Model loading and export: Add a reproducible example with pinned checkpoint and dependency revisions, a PyTorch reference, and backend selection. Export the neural computation to .pte plus any required backend artifacts. Define input/output shapes, supported dtypes, and bounded sequence lengths or padded chunk sizes. Keep the shared model representation portable across the four backends.
• Native inference: Provide a C++ runner and audio-file CLI covering mel preprocessing, chunk scheduling, AOSC/FIFO updates, and conversion of activity probabilities into speaker/start/end segments. Expose streaming feed, reset, and final-flush behavior; document which work runs on the host. Inference should run without Python or NeMo installed.
• Streaming and offline modes: Support the published 30.4 s offline-style configuration and 1.04 s, 0.64 s, and 0.32 s streaming presets. Preserve arrival-order speaker labels, look-ahead handling, padding masks, and timestamp alignment across chunks. These values are input-buffer latency, excluding compute time.

Backend workstreams

Backend	Target	Status	Work to validate and enable
MLX	Apple Silicon GPU	Initial support landed in #23137; XNNPACK/CPU preprocessing and MLX neural inference.	Lower the encoder and speaker head through the MLX delegate; validate RoPE, masked attention, Conv1D, and varying chunk/cache lengths.
CUDA	NVIDIA GPU	Open.	Integrate the CUDA export/runtime path; validate attention, shape handling, supported precision, and host/device transfer costs. Document the tested GPU and CUDA requirements.
XNNPACK	Arm and x86 CPU	Preprocessor enabled; neural inference remains open.	Establish a floating-point CPU baseline; inspect attention/matmul, normalization, and Conv1D delegation, and tune threading. Validate dynamic shapes or provide padded/static variants where needed.
Vulkan	Supported Android and desktop GPUs	Open.	Validate attention, RoPE, Conv1D, tensor layouts, shape changes, and device limits. Add lowering/kernel support as needed and document any partitions assigned to XNNPACK or portable kernels during export.

These are validation targets, not confirmed operator gaps. Record actual delegation coverage and remaining blockers for each backend. Establish floating-point correctness first; evaluate reduced precision and quantization separately against that baseline.

Acceptance criteria

• Each backend has documented export/build/run commands and a successful end-to-end run on named hardware, with delegation coverage and fallback operators reported.
• Compare preprocessing, per-frame probabilities, and final segments against the same pinned upstream reference. Report diarization error rate (DER) with the dataset, scoring settings, and agreed numerical/quality tolerances.
• Cover silence, overlapping speech, speaker arrivals, up to eight speakers, short/final chunks, reset between recordings, cache rollover, and long recordings with bounded streaming memory.
• Report model/artifact size, peak memory, real-time factor, and p50/p95 chunk processing latency at batch size 1, including preprocessing and state updates. Separate initialization/warm-up, input buffering, and steady-state compute; record hardware, precision, thread count, and streaming preset.
• Add regression coverage for export, runtime correctness, and streaming state handling, with documentation of supported configurations and remaining limitations.

Alternatives

The upstream NeMo and Transformers implementations provide reference inference. NeMo-Speech.cpp provides another native deployment option. This request brings the model into ExecuTorch's runtime and delegate ecosystem.

Additional context

• Model card and recommended streaming configurations
• Checkpoint configuration
• Transformers integration

cc @SS-JIA @manuelcandales @digantdesai @cbilgin @GregoryComer @JakeStevens @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng @Gasoonjia @metascroy

[mergennachin]

#23152 - XNNPACK
#23155 - Vulkan

[mergennachin]

#23167 -CUDA