@@ -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, XNNPACK)
+# - nemotron3-diarization: Speaker diarization model (MLX, XNNPACK, CUDA)
 # - supertonic: Text-to-speech model (MLX)
 # - silero_vad: Voice activity detection model (CPU)
 # - llama:    Text generation model (CPU)
@@ -124,6 +124,7 @@ help:
 	@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 "  nemotron3-diarization-cuda - Build Nemotron 3 Diarization runner with CUDA 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"
@@ -299,6 +300,11 @@ nemotron3-diarization-cpu:
 	cmake --workflow --preset llm-release
 	cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-cpu
 
+.PHONY: nemotron3-diarization-cuda
+nemotron3-diarization-cuda:
+	cmake --workflow --preset llm-release-cuda
+	cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-cuda
+
 sortformer-cuda:
 	@echo "==> Building and installing ExecuTorch with CUDA..."
 	cmake --workflow --preset llm-release-cuda
 
@@ -42,6 +42,13 @@ if(TARGET kleidiai)
   target_link_libraries(nemotron3_diarization PUBLIC kleidiai)
 endif()
 executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
+if(EXECUTORCH_BUILD_CUDA)
+  find_package(CUDAToolkit REQUIRED)
+  target_link_libraries(nemotron3_diarization PUBLIC aoti_cuda_backend)
+  if(NOT MSVC)
+    executorch_target_link_options_shared_lib(aoti_cuda_backend)
+  endif()
+endif()
 add_executable(nemotron3_diarization_runner main.cpp)
 target_include_directories(
   nemotron3_diarization_runner
@@ -50,6 +57,11 @@ target_include_directories(
 target_link_libraries(
   nemotron3_diarization_runner PRIVATE nemotron3_diarization gflags
 )
+if(EXECUTORCH_BUILD_CUDA)
+  target_compile_definitions(
+    nemotron3_diarization_runner PRIVATE EXECUTORCH_BUILD_CUDA
+  )
+endif()
 if(EXECUTORCH_BUILD_MLX)
   if(NOT TARGET mlxdelegate OR NOT TARGET mlx)
     message(
@@ -63,3 +75,12 @@ if(EXECUTORCH_BUILD_MLX)
   )
   executorch_target_copy_mlx_metallib(nemotron3_diarization_runner)
 endif()
+if(MSVC AND EXECUTORCH_BUILD_CUDA)
+  add_custom_command(
+    TARGET nemotron3_diarization_runner
+    POST_BUILD
+    COMMAND
+      ${CMAKE_COMMAND} -E copy_if_different $<TARGET_FILE:aoti_cuda_shims>
+      $<TARGET_FILE_DIR:nemotron3_diarization_runner>
+  )
+endif()
@@ -6,6 +6,7 @@
       "hidden": true,
       "binaryDir": "${sourceDir}/../../../cmake-out/examples/models/nemotron3_diarization",
       "cacheVariables": {
+        "EXECUTORCH_BUILD_CUDA": "OFF",
         "EXECUTORCH_BUILD_MLX": "OFF",
         "CMAKE_BUILD_TYPE": "Release",
         "CMAKE_FIND_ROOT_PATH": "${sourceDir}/../../../cmake-out",
@@ -28,6 +29,21 @@
         "rhs": "Darwin"
       }
     },
+    {
+      "name": "nemotron3-diarization-cuda",
+      "displayName": "Nemotron 3 Diarization runner (CUDA)",
+      "inherits": [
+        "nemotron3-diarization-base"
+      ],
+      "cacheVariables": {
+        "EXECUTORCH_BUILD_CUDA": "ON"
+      },
+      "condition": {
+        "type": "inList",
+        "string": "${hostSystemName}",
+        "list": ["Linux", "Windows"]
+      }
+    },
     {
       "name": "nemotron3-diarization-cpu",
       "displayName": "Nemotron 3 Diarization runner (XNNPACK)",
@@ -37,6 +53,11 @@
     }
   ],
   "buildPresets": [
+    {
+      "name": "nemotron3-diarization-cuda",
+      "configurePreset": "nemotron3-diarization-cuda",
+      "configuration": "Release"
+    },
     {
       "name": "nemotron3-diarization-mlx",
       "configurePreset": "nemotron3-diarization-mlx",
@@ -49,6 +70,19 @@
     }
   ],
   "workflowPresets": [
+    {
+      "name": "nemotron3-diarization-cuda",
+      "steps": [
+        {
+          "type": "configure",
+          "name": "nemotron3-diarization-cuda"
+        },
+        {
+          "type": "build",
+          "name": "nemotron3-diarization-cuda"
+        }
+      ]
+    },
     {
       "name": "nemotron3-diarization-mlx",
       "steps": [
 
@@ -4,9 +4,9 @@
 # From the repository root:
 python -m pip install git+https://github.com/huggingface/transformers "librosa>=0.10"
 
-# Choose mlx (Apple Silicon) or xnnpack (CPU).
+# Choose mlx (Apple Silicon), xnnpack (CPU), or cuda (NVIDIA GPU).
 BACKEND=mlx
-# Export: BF16 for MLX/XNNPACK. Add --dtype fp32 to use FP32.
+# Export: BF16 by default. Add --dtype fp32 to use FP32.
 python -m executorch.examples.models.nemotron3_diarization.export_nemotron \
   --hf-model nvidia/Nemotron-3-Diarization \
   --backend "$BACKEND" \
@@ -17,9 +17,20 @@ 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
+DATA_PATH=""
+if [ "$BACKEND" = cuda ]; then
+  DATA_PATH="nemotron_exports/cuda/aoti_cuda_blob.ptd"
+fi
 cmake-out/examples/models/nemotron3_diarization/nemotron3_diarization_runner \
   --model_path="nemotron_exports/$BACKEND/nemotron3_diarization.pte" \
+  --data_path="$DATA_PATH" \
   --audio_path=audio_16khz_mono.wav \
   --preset=offline \
   --output=segments.json
 ```
+
+CUDA export requires a CUDA-enabled PyTorch installation and a CUDA toolkit.
+The preprocessor runs on CPU with XNNPACK; `pre_encode` and `encode` run on CUDA,
+with float32 inputs and outputs and BF16 (or FP32) model computation. Export
+produces `nemotron3_diarization.pte` and `aoti_cuda_blob.ptd`; pass the latter
+through `--data_path` as shown above.
@@ -63,22 +63,26 @@ def forward(self, features):
 class Encode(nn.Module):
     """Adapt the stateless Transformers network to the native streaming ABI."""
 
-    def __init__(self, model):
+    def __init__(self, model, *, expand_attention_mask=False):
         super().__init__()
         self.model = model.model
         self.classifier = model.classifier
         self.dtype = model.dtype
+        self.expand_attention_mask = expand_attention_mask
 
     def forward(self, embeddings, lengths):
         positions = torch.arange(embeddings.shape[1], device=embeddings.device)
         mask = (positions[None, :] < lengths[:, None])[:, None, None, :]
+        if self.expand_attention_mask:
+            # CUDA SDPA requires the full query dimension in the mask.
+            mask = mask.expand(-1, -1, embeddings.shape[1], -1)
         encoded = self.model(
             inputs_embeds=embeddings.to(self.dtype), attention_mask=mask
         )
         return self.classifier(encoded.last_hidden_state).sigmoid().float()
 
 
-def capture_model(model, feature_extractor):
+def capture_model(model, feature_extractor, *, expand_attention_mask=False):
     """Capture backend-independent programs; streaming state stays in the runner."""
     audio = model.config.audio_config
     head = model.config.head_config
@@ -145,7 +149,7 @@ def capture_model(model, feature_extractor):
                 strict=False,
             ),
             "encode": export(
-                Encode(model),
+                Encode(model, expand_attention_mask=expand_attention_mask),
                 (
                     torch.zeros(1, 128, audio.hidden_size),
                     torch.tensor([128], dtype=torch.int64),
@@ -169,27 +173,49 @@ def export_model(
         from executorch.backends.mlx.partitioner import MLXPartitioner
         from executorch.backends.mlx.passes import get_default_passes
 
-        partitioners = [MLXPartitioner()]
+        partitioners = {name: [MLXPartitioner()] for name in ("pre_encode", "encode")}
         passes = get_default_passes()
     elif backend == "xnnpack":
-        partitioners = [XnnpackPartitioner(enable_bf16=dtype == torch.bfloat16)]
+        partitioners = {
+            name: [XnnpackPartitioner(enable_bf16=dtype == torch.bfloat16)]
+            for name in ("pre_encode", "encode")
+        }
+    elif backend == "cuda":
+        from executorch.backends.cuda.cuda_backend import CudaBackend
+        from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+        from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+        # The custom Triton SDPA kernel supports BF16 only.
+        triton_mode = b"ON" if dtype == torch.bfloat16 else b"OFF"
+        partitioners = {
+            name: [
+                CudaPartitioner(
+                    [
+                        CudaBackend.generate_method_name_compile_spec(name),
+                        CompileSpec("triton_kernel_mode", triton_mode),
+                    ]
+                )
+            ]
+            for name in ("pre_encode", "encode")
+        }
     else:
         raise ValueError(f"Unknown backend: {backend}")
+    partitioners["preprocessor"] = [XnnpackPartitioner()]
 
     model = Nemotron3DiarizationForAudioFrameClassification.from_pretrained(
         model_id, revision=revision, dtype=dtype
     )
     processor = Nemotron3DiarizationProcessor.from_pretrained(
         model_id, revision=revision
     )
-    programs, metadata = capture_model(model, processor.feature_extractor)
+    programs, metadata = capture_model(
+        model,
+        processor.feature_extractor,
+        expand_attention_mask=backend == "cuda",
+    )
     edge = to_edge_transform_and_lower(
         programs,
-        partitioner={
-            "preprocessor": [XnnpackPartitioner()],
-            "pre_encode": partitioners,
-            "encode": partitioners,
-        },
+        partitioner=partitioners,
         transform_passes={
             "preprocessor": [],
             "pre_encode": passes,
@@ -212,6 +238,9 @@ def export_model(
     with path.open("wb") as output:
         program.write_to_file(output)
     print(f"Saved {path} ({path.stat().st_size / 2**20:.1f} MiB)")
+    if backend == "cuda":
+        program.write_tensor_data_to_file(str(output_dir))
+        print(f"Saved {output_dir / 'aoti_cuda_blob.ptd'}")
     return path
 
 
@@ -223,7 +252,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("--backend", choices=("mlx", "xnnpack", "cuda"), default="mlx")
     parser.add_argument("--output-dir", type=Path, default=Path("nemotron_exports"))
     parser.add_argument(
         "--dtype",
 
@@ -27,6 +27,7 @@ DEFINE_string(
     model_path,
     "nemotron_exports/nemotron3_diarization.pte",
     "Exported model");
+DEFINE_string(data_path, "", "External CUDA data file (aoti_cuda_blob.ptd)");
 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");
@@ -46,6 +47,12 @@ int main(int argc, char** argv) {
       throw std::invalid_argument(
           "Provide --audio_path, positive --feed_samples, and --threshold in [0,1]");
     }
+#ifdef EXECUTORCH_BUILD_CUDA
+    if (FLAGS_data_path.empty()) {
+      throw std::invalid_argument(
+          "Provide --data_path pointing to aoti_cuda_blob.ptd for CUDA");
+    }
+#endif
 #ifdef EXECUTORCH_BUILD_MLX
     if (FLAGS_mlx_cache_limit_mb < 0 || FLAGS_mlx_memory_limit_mb <= 0) {
       throw std::invalid_argument(
@@ -58,7 +65,8 @@ int main(int argc, char** argv) {
 #endif
     nemotron3::Runner runner(
         FLAGS_model_path,
-        nemotron3::StreamingConfig::from_preset(FLAGS_preset));
+        nemotron3::StreamingConfig::from_preset(FLAGS_preset),
+        FLAGS_data_path);
     auto header = executorch::extension::llm::load_wav_header(FLAGS_audio_path);
     if (!header || header->NumOfChan != 1 ||
         header->SamplesPerSec != runner.sample_rate()) {
 
@@ -93,8 +93,14 @@ StreamingConfig StreamingConfig::from_preset(const std::string& preset) {
   throw std::invalid_argument("Unknown streaming preset: " + preset);
 }
 
-Runner::Runner(const std::string& model_path, StreamingConfig config)
-    : model_(model_path, executorch::extension::Module::LoadMode::Mmap),
+Runner::Runner(
+    const std::string& model_path,
+    StreamingConfig config,
+    const std::string& data_path)
+    : model_(
+          model_path,
+          data_path,
+          executorch::extension::Module::LoadMode::Mmap),
       config_(config) {
   if (model_.load() != Error::Ok) {
     throw std::runtime_error("Could not load " + model_path);
 
@@ -27,7 +27,10 @@ struct StreamingConfig {
 
 class Runner {
  public:
-  explicit Runner(const std::string& model_path, StreamingConfig config = {});
+  explicit Runner(
+      const std::string& model_path,
+      StreamingConfig config = {},
+      const std::string& data_path = "");
 
   // Returns new 10 ms frames, with num_speakers() probabilities per frame.
   // Audio is mono at sample_rate(). final=true flushes the remaining lookahead.