[examples] Enable Nemotron 3 Diarization on XNNPACK (pytorch#23152)

### Summary

Add XNNPACK CPU export and native runner build targets for the
Transformers-based Nemotron 3 Diarization example. The CPU path defaults
to BF16, with `--dtype fp32` available explicitly, and reuses the
existing streaming runner and XNNPACK preprocessor. MLX-specific code
and linking are enabled only for MLX builds.

Rebased onto main after pytorch#23153 landed. This PR contains only the six
example/build files; the CPU kernel fixes and tests are already in main.
Vulkan support is separate in pytorch#23155.

### Test plan

Built the native CPU runner using `make nemotron3-diarization-cpu` on
macOS arm64, with the CPU fixes from main. Verified the runner links
XNNPACK without either the MLX or Vulkan delegate. Fresh exports pass
for the default BF16 mode and explicit FP32 mode.

All eight production CLI runs (two dtypes across offline, low, very_low,
and ultra_low) retain the previous segments. All eight streaming
integration runs pass partitioning, reset, short final-window, and
input-validation checks. Also exercised 11 method shapes per dtype and
compared speech, partial-tail, and silence inputs across all four
presets against saved Transformers references.

The 12 FP32 audio comparisons pass at `rtol=1e-4, atol=1e-4`, with
maximum probability error 1.3709068298339844e-6. The 12 BF16 audio
comparisons have maximum probability error 0.015625; four of 52,896
values exceed `rtol=1e-2, atol=5e-3`. Both dtypes have zero
speaker-decision disagreements. BF16 is approximate and is not claimed
to pass that tight numerical tolerance.

Export CLI, CMake presets, `lintrunner`, and `git diff --check` pass.
The rebase preserves the example patch unchanged and removes the CPU
commits already merged through pytorch#23153.

Supplemental evaluation scripts and artifacts remain local.

Authored with AI assistance from OpenAI Codex.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Author: mergennachin

Makefile

@@ -19,7 +19,7 @@
 # - whisper:  Speech recognition model (CPU, CUDA, Metal)
 # - parakeet: Speech recognition model (CPU, CUDA, Metal, MLX)
 # - sortformer: Speaker diarization model (CPU, CUDA)
-# - nemotron3-diarization: Speaker diarization model (MLX)
+# - nemotron3-diarization: Speaker diarization model (MLX, XNNPACK)
 # - supertonic: Text-to-speech model (MLX)
 # - silero_vad: Voice activity detection model (CPU)
 # - llama:    Text generation model (CPU)
@@ -123,6 +123,7 @@ help:
 	@echo "  sortformer-cuda     - Build Sortformer runner with CUDA backend"
 	@echo "  sortformer-cpu      - Build Sortformer runner with CPU backend"
 	@echo "  nemotron3-diarization-mlx - Build Nemotron 3 Diarization runner with MLX backend"
+	@echo "  nemotron3-diarization-cpu - Build Nemotron 3 Diarization runner with XNNPACK backend"
 	@echo "  supertonic-mlx      - Build Supertonic runner with MLX backend"
 	@echo "  silero-vad-cpu      - Build Silero VAD runner with CPU backend"
 	@echo "  llama-cuda          - Build Llama runner with CUDA backend"
@@ -293,6 +294,11 @@ nemotron3-diarization-mlx:
 	cmake --build --preset mlx-release-install
 	cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-mlx
 
+.PHONY: nemotron3-diarization-cpu
+nemotron3-diarization-cpu:
+	cmake --workflow --preset llm-release
+	cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-cpu
+
 sortformer-cuda:
 	@echo "==> Building and installing ExecuTorch with CUDA..."
 	cmake --workflow --preset llm-release-cuda

examples/models/nemotron3_diarization/CMakeLists.txt

@@ -15,16 +15,11 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
 set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags)
 find_package(gflags REQUIRED)
 find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH)
-if(NOT TARGET mlxdelegate
-   OR NOT TARGET mlx
-   OR NOT TARGET xnnpack_backend
-)
+if(NOT TARGET xnnpack_backend)
   message(
-    FATAL_ERROR
-      "Build and install ExecuTorch with EXECUTORCH_BUILD_MLX=ON and EXECUTORCH_BUILD_XNNPACK=ON"
+    FATAL_ERROR "Build and install ExecuTorch with EXECUTORCH_BUILD_XNNPACK=ON"
   )
 endif()
-executorch_target_link_options_shared_lib(mlxdelegate)
 executorch_target_link_options_shared_lib(xnnpack_backend)
 
 add_library(nemotron3_diarization STATIC runner.cpp)
@@ -53,7 +48,18 @@ target_include_directories(
   PRIVATE ${EXECUTORCH_ROOT}/third-party/json/include
 )
 target_link_libraries(
-  nemotron3_diarization_runner PRIVATE nemotron3_diarization gflags mlxdelegate
-                                       mlx
+  nemotron3_diarization_runner PRIVATE nemotron3_diarization gflags
 )
-executorch_target_copy_mlx_metallib(nemotron3_diarization_runner)
+if(EXECUTORCH_BUILD_MLX)
+  if(NOT TARGET mlxdelegate OR NOT TARGET mlx)
+    message(
+      FATAL_ERROR "Build and install ExecuTorch with EXECUTORCH_BUILD_MLX=ON"
+    )
+  endif()
+  executorch_target_link_options_shared_lib(mlxdelegate)
+  target_link_libraries(nemotron3_diarization_runner PRIVATE mlxdelegate mlx)
+  target_compile_definitions(
+    nemotron3_diarization_runner PRIVATE EXECUTORCH_BUILD_MLX
+  )
+  executorch_target_copy_mlx_metallib(nemotron3_diarization_runner)
+endif()

examples/models/nemotron3_diarization/CMakePresets.json

@@ -2,32 +2,77 @@
   "version": 6,
   "configurePresets": [
     {
-      "name": "nemotron3-diarization-mlx",
-      "displayName": "Nemotron 3 Diarization runner (MLX)",
+      "name": "nemotron3-diarization-base",
+      "hidden": true,
       "binaryDir": "${sourceDir}/../../../cmake-out/examples/models/nemotron3_diarization",
       "cacheVariables": {
-        "EXECUTORCH_BUILD_MLX": "ON",
+        "EXECUTORCH_BUILD_MLX": "OFF",
         "CMAKE_BUILD_TYPE": "Release",
-        "CMAKE_OSX_DEPLOYMENT_TARGET": "14.0",
         "CMAKE_FIND_ROOT_PATH": "${sourceDir}/../../../cmake-out",
         "CMAKE_PREFIX_PATH": "${sourceDir}/../../../cmake-out"
+      }
+    },
+    {
+      "name": "nemotron3-diarization-mlx",
+      "displayName": "Nemotron 3 Diarization runner (MLX)",
+      "inherits": [
+        "nemotron3-diarization-base"
+      ],
+      "cacheVariables": {
+        "EXECUTORCH_BUILD_MLX": "ON",
+        "CMAKE_OSX_DEPLOYMENT_TARGET": "14.0"
       },
-      "condition": {"type": "equals", "lhs": "${hostSystemName}", "rhs": "Darwin"}
+      "condition": {
+        "type": "equals",
+        "lhs": "${hostSystemName}",
+        "rhs": "Darwin"
+      }
+    },
+    {
+      "name": "nemotron3-diarization-cpu",
+      "displayName": "Nemotron 3 Diarization runner (XNNPACK)",
+      "inherits": [
+        "nemotron3-diarization-base"
+      ]
     }
   ],
   "buildPresets": [
     {
       "name": "nemotron3-diarization-mlx",
       "configurePreset": "nemotron3-diarization-mlx",
       "configuration": "Release"
+    },
+    {
+      "name": "nemotron3-diarization-cpu",
+      "configurePreset": "nemotron3-diarization-cpu",
+      "configuration": "Release"
     }
   ],
   "workflowPresets": [
     {
       "name": "nemotron3-diarization-mlx",
       "steps": [
-        {"type": "configure", "name": "nemotron3-diarization-mlx"},
-        {"type": "build", "name": "nemotron3-diarization-mlx"}
+        {
+          "type": "configure",
+          "name": "nemotron3-diarization-mlx"
+        },
+        {
+          "type": "build",
+          "name": "nemotron3-diarization-mlx"
+        }
+      ]
+    },
+    {
+      "name": "nemotron3-diarization-cpu",
+      "steps": [
+        {
+          "type": "configure",
+          "name": "nemotron3-diarization-cpu"
+        },
+        {
+          "type": "build",
+          "name": "nemotron3-diarization-cpu"
+        }
       ]
     }
   ]

examples/models/nemotron3_diarization/README.md

@@ -1,22 +1,24 @@
-# Nemotron 3 Diarization on MLX
+# Nemotron 3 Diarization
 
 ```bash
-# From the repository root on Apple Silicon:
+# From the repository root:
 python -m pip install git+https://github.com/huggingface/transformers "librosa>=0.10"
 
-# Export (approximate BF16; add --dtype fp32 for comparison with Transformers)
-# Uses the Transformers frontend.
+# Choose mlx (Apple Silicon) or xnnpack (CPU).
+BACKEND=mlx
+# Export: BF16 for MLX/XNNPACK. Add --dtype fp32 to use FP32.
 python -m executorch.examples.models.nemotron3_diarization.export_nemotron \
   --hf-model nvidia/Nemotron-3-Diarization \
-  --output-dir nemotron_exports
+  --backend "$BACKEND" \
+  --output-dir "nemotron_exports/$BACKEND"
 
-# Build
-make nemotron3-diarization-mlx
+# Build the matching runner.
+make "nemotron3-diarization-${BACKEND/xnnpack/cpu}"
 
 # Convert to mono 16 kHz PCM16 and run
 ffmpeg -i audio.wav -ac 1 -ar 16000 -c:a pcm_s16le audio_16khz_mono.wav
 cmake-out/examples/models/nemotron3_diarization/nemotron3_diarization_runner \
-  --model_path=nemotron_exports/nemotron3_diarization.pte \
+  --model_path="nemotron_exports/$BACKEND/nemotron3_diarization.pte" \
   --audio_path=audio_16khz_mono.wav \
   --preset=offline \
   --output=segments.json

examples/models/nemotron3_diarization/export_nemotron.py

@@ -157,13 +157,25 @@ def capture_model(model, feature_extractor):
     return programs, metadata
 
 
-def export_model(model_id, output_dir, revision=None, dtype=torch.bfloat16):
-    from executorch.backends.mlx.partitioner import MLXPartitioner
-    from executorch.backends.mlx.passes import get_default_passes
+def export_model(
+    model_id, output_dir, revision=None, dtype=torch.bfloat16, backend="mlx"
+):
     from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
         XnnpackPartitioner,
     )
 
+    passes = []
+    if backend == "mlx":
+        from executorch.backends.mlx.partitioner import MLXPartitioner
+        from executorch.backends.mlx.passes import get_default_passes
+
+        partitioners = [MLXPartitioner()]
+        passes = get_default_passes()
+    elif backend == "xnnpack":
+        partitioners = [XnnpackPartitioner(enable_bf16=dtype == torch.bfloat16)]
+    else:
+        raise ValueError(f"Unknown backend: {backend}")
+
     model = Nemotron3DiarizationForAudioFrameClassification.from_pretrained(
         model_id, revision=revision, dtype=dtype
     )
@@ -175,13 +187,13 @@ def export_model(model_id, output_dir, revision=None, dtype=torch.bfloat16):
         programs,
         partitioner={
             "preprocessor": [XnnpackPartitioner()],
-            "pre_encode": [MLXPartitioner()],
-            "encode": [MLXPartitioner()],
+            "pre_encode": partitioners,
+            "encode": partitioners,
         },
         transform_passes={
             "preprocessor": [],
-            "pre_encode": get_default_passes(),
-            "encode": get_default_passes(),
+            "pre_encode": passes,
+            "encode": passes,
         },
         constant_methods=metadata,
         compile_config=EdgeCompileConfig(
@@ -211,6 +223,7 @@ def main():
         help="Hugging Face model ID or local Transformers checkpoint directory",
     )
     parser.add_argument("--revision", help="Hugging Face model revision")
+    parser.add_argument("--backend", choices=("mlx", "xnnpack"), default="mlx")
     parser.add_argument("--output-dir", type=Path, default=Path("nemotron_exports"))
     parser.add_argument(
         "--dtype",
@@ -220,7 +233,7 @@ def main():
     )
     args = parser.parse_args()
     dtype = {"bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype]
-    export_model(args.hf_model, args.output_dir, args.revision, dtype)
+    export_model(args.hf_model, args.output_dir, args.revision, dtype, args.backend)
 
 
 if __name__ == "__main__":

examples/models/nemotron3_diarization/main.cpp

@@ -15,9 +15,12 @@
 #include <stdexcept>
 
 #include <gflags/gflags.h>
-#include <mlx/memory.h>
 #include <nlohmann/json.hpp>
 
+#ifdef EXECUTORCH_BUILD_MLX
+#include <mlx/memory.h>
+#endif
+
 #include <executorch/extension/llm/runner/wav_loader.h>
 
 DEFINE_string(
@@ -28,25 +31,31 @@ DEFINE_string(audio_path, "", "Mono 16 kHz WAV file (PCM16 or float32)");
 DEFINE_string(preset, "offline", "offline, low, very_low, or ultra_low");
 DEFINE_int32(feed_samples, 4096, "Samples per streaming feed");
 DEFINE_double(threshold, 0.5, "Speaker activity threshold");
+#ifdef EXECUTORCH_BUILD_MLX
 DEFINE_int32(mlx_cache_limit_mb, 128, "Maximum unused MLX buffer cache in MiB");
 DEFINE_int32(mlx_memory_limit_mb, 4096, "MLX working memory guideline in MiB");
+#endif
 DEFINE_string(output, "", "Optional JSON output file");
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
   try {
     if (FLAGS_audio_path.empty() || FLAGS_feed_samples <= 0 ||
-        FLAGS_mlx_cache_limit_mb < 0 || FLAGS_mlx_memory_limit_mb <= 0 ||
         !std::isfinite(FLAGS_threshold) || FLAGS_threshold < 0 ||
         FLAGS_threshold > 1) {
       throw std::invalid_argument(
-          "Provide --audio_path, positive --feed_samples and --mlx_memory_limit_mb, "
-          "nonnegative --mlx_cache_limit_mb, and --threshold in [0,1]");
+          "Provide --audio_path, positive --feed_samples, and --threshold in [0,1]");
+    }
+#ifdef EXECUTORCH_BUILD_MLX
+    if (FLAGS_mlx_cache_limit_mb < 0 || FLAGS_mlx_memory_limit_mb <= 0) {
+      throw std::invalid_argument(
+          "Provide nonnegative --mlx_cache_limit_mb and positive --mlx_memory_limit_mb");
     }
     mlx::core::set_cache_limit(
         static_cast<size_t>(FLAGS_mlx_cache_limit_mb) << 20);
     mlx::core::set_memory_limit(
         static_cast<size_t>(FLAGS_mlx_memory_limit_mb) << 20);
+#endif
     nemotron3::Runner runner(
         FLAGS_model_path,
         nemotron3::StreamingConfig::from_preset(FLAGS_preset));