[executorch][cuda] Deprecate triton_kernel_mode=OFF

backends/cuda/cuda_backend.py

@@ -15,6 +15,7 @@
 import sys
 import threading
 import typing
+import warnings
 import weakref
 from importlib import resources
 from typing import Any, Dict, final, List, Optional
@@ -1253,23 +1254,29 @@ def get_decomposition_table(cls) -> Dict[Any, Any]:
     @classmethod
     def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]:
         """
-        Return CUDA-specific passes: ReplaceEdgeOpWithTritonOpPass.
+        Return CUDA-specific passes, ending with ReplaceEdgeOpWithTritonOpPass.
 
-        The Triton kernel replacement behavior can be controlled via compile_specs:
-        - triton_kernel_mode="ON": Use Triton kernels for the calls they support; other calls use the regular lowering
-        - triton_kernel_mode="OFF": Never use Triton kernels and fallback to other implementations like cuda or decomposed operator.
+        The replacement pass uses Triton kernels for the calls they support and
+        leaves other calls to the regular lowering. The triton_kernel_mode
+        compile spec is deprecated: "OFF" is treated as "ON".
         """
-        # Parse compile_specs for triton_kernel_mode
-
-        triton_kernel_mode = "ON"  # Default mode
         for spec in compile_specs:
             if spec.key == "triton_kernel_mode":
                 mode = spec.value.decode("utf-8").upper()
                 if mode not in ["ON", "OFF"]:
                     raise ValueError(
                         f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
-                triton_kernel_mode = mode
+                if mode == "OFF":
+                    warnings.warn(
+                        "triton_kernel_mode=OFF is deprecated and is now treated as "
+                        "ON: the CUDA backend replaces the calls its Triton kernels "
+                        "support and leaves the others to the regular lowering. "
+                        "Remove the triton_kernel_mode compile spec; it will be "
+                        "removed in a future release.",
+                        FutureWarning,
+                        stacklevel=2,
+                    )
         passes = [
             RemoveAssertAsyncPass(),
             ReplaceComplexWithRealPass(),
@@ -1294,8 +1301,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                 passes.append(
                     LowerOffGraphKVPass(parse_offgraph_kv_manifest(spec.value))
                 )
-        if triton_kernel_mode == "ON":
-            passes.append(ReplaceEdgeOpWithTritonOpPass())
+        passes.append(ReplaceEdgeOpWithTritonOpPass())
         return passes
 
     @classmethod

backends/cuda/tests/targets.bzl

@@ -19,6 +19,7 @@ def define_common_targets(is_fbcode = False):
             "//caffe2:torch",
             "//executorch/backends/cuda:cuda_backend",
             "//executorch/backends/cuda:cuda_partitioner",
+            "//executorch/backends/cuda:triton_replacement_pass",
             "//executorch/exir:lib",
             "//executorch/exir/backend:backend_api",
             "//executorch/exir/backend:compile_spec_schema",

backends/cuda/tests/test_cuda_export.py

@@ -18,6 +18,9 @@
 from executorch.backends.cuda.cuda_backend import CudaBackend
 from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
 from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass
+from executorch.backends.cuda.triton.replacement_pass import (
+    ReplaceEdgeOpWithTritonOpPass,
+)
 from executorch.examples.models.toy_model import SdpaModule
 from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower
 from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -941,29 +944,16 @@ def forward(self, *qkv):
                 )
 
     def test_triton_kernel_mode_off(self):
-        """
-        Test CUDA export with triton_kernel_mode set to OFF for SDPA kernel.
-        This validates that the backend correctly processes the triton_kernel_mode
-        compile spec and can export SDPA operations without Triton kernel replacements.
-        When triton_kernel_mode is OFF, SDPA should be decomposed using the MATH backend.
-        """
-
-        sdpa = SdpaModule()
-
-        # Create compile specs with triton_kernel_mode set to OFF
-        compile_specs = [
-            CudaBackend.generate_method_name_compile_spec("forward"),
-            CompileSpec(key="triton_kernel_mode", value=b"OFF"),
-        ]
-
-        # Test export with triton_kernel_mode=OFF
-        edge_program_manager = self._export_to_cuda_with_lower(
-            sdpa.get_eager_model(), sdpa.get_example_inputs(), compile_specs
-        )
-        self.assertIsNotNone(
-            edge_program_manager,
-            "SDPA kernel export with triton_kernel_mode=OFF failed",
-        )
+        """The deprecated OFF mode warns and is treated as ON."""
+        on = CudaBackend.get_custom_passes([CompileSpec("triton_kernel_mode", b"ON")])
+        with self.assertWarnsRegex(
+            FutureWarning, "triton_kernel_mode=OFF is deprecated"
+        ):
+            off = CudaBackend.get_custom_passes(
+                [CompileSpec("triton_kernel_mode", b"OFF")]
+            )
+        self.assertEqual([type(p) for p in off], [type(p) for p in on])
+        self.assertIsInstance(off[-1], ReplaceEdgeOpWithTritonOpPass)
 
     def test_missing_scatter_kernel_fails_at_export(self):
         """The scatter fallback writes its kernel call itself, and the check for

backends/cuda/tests/test_output_stride_rearrange.py

@@ -99,7 +99,7 @@ def _run_rocm_runner(module, inputs, output_strides, pte, ptd, tmpdir):
     ]
 
 
-def _run_et_aoti(module, inputs, triton_mode="OFF"):
+def _run_et_aoti(module, inputs):
     """Export and run a model through the ExecuTorch CUDA AOTI backend."""
     from executorch.backends.cuda.cuda_backend import CudaBackend
     from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
@@ -108,14 +108,10 @@ def _run_et_aoti(module, inputs, triton_mode="OFF"):
         ExecutorchBackendConfig,
         to_edge_transform_and_lower,
     )
-    from executorch.exir.backend.compile_spec_schema import CompileSpec
     from executorch.exir.passes import MemoryPlanningPass
 
     ep = torch.export.export(module, inputs, strict=True)
-    compile_specs = [
-        CudaBackend.generate_method_name_compile_spec("forward"),
-        CompileSpec(key="triton_kernel_mode", value=triton_mode.encode()),
-    ]
+    compile_specs = [CudaBackend.generate_method_name_compile_spec("forward")]
     partitioner = {"forward": [CudaPartitioner(compile_specs)]}
     et_prog = to_edge_transform_and_lower(
         {"forward": ep},
@@ -170,19 +166,15 @@ def setUpClass(cls):
             raise unittest.SkipTest("BF16 not supported on this GPU")
 
     def _check_sdpa(self, module, inputs, label):
-        """Run model through ExecuTorch AOTI with both Triton ON and OFF."""
+        """Run model through ExecuTorch AOTI and compare with eager."""
         with torch.no_grad():
             eager = module(*inputs).float().cpu()
 
-        for mode in ["OFF", "ON"]:
-            with self.subTest(triton_mode=mode):
-                result = _run_et_aoti(module, inputs, triton_mode=mode)[0].float()
-                rel = (result - eager).abs() / eager.abs().clamp(min=1e-6)
-                self.assertLess(
-                    rel.mean().item(),
-                    0.2,
-                    f"{label} triton={mode} mean_rel={rel.mean():.4f} too large",
-                )
+        result = _run_et_aoti(module, inputs)[0].float()
+        rel = (result - eager).abs() / eager.abs().clamp(min=1e-6)
+        self.assertLess(
+            rel.mean().item(), 0.2, f"{label} mean_rel={rel.mean():.4f} too large"
+        )
 
     def test_fast_path_no_mask(self):
         """No mask SDPA — output strides match, uses fast byte copy path."""

examples/models/nemotron3_diarization/export_nemotron.py

@@ -207,18 +207,10 @@ def export_model(model_id, output_dir, revision=None, dtype=None, backend="mlx")
     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),
-                    ]
-                )
+                CudaPartitioner([CudaBackend.generate_method_name_compile_spec(name)])
             ]
             for name in ("pre_encode", "encode")
         }

examples/models/voxtral_tts/export_voxtral_tts.py

@@ -143,7 +143,8 @@ def _export_codec_pte(model, args, device: str) -> None:
     Codec convs are expressed as unfold + matmul / matmul + Fold
     (model.py:_conv1d_as_matmul / _conv_transpose1d_as_matmul) so AOTI's CUDA
     backend can lower them via Triton mm kernels. CodecAttention uses an
-    additive ALiBi mask which is fine for ATen SDPA when triton_kernel_mode=OFF.
+    additive ALiBi mask, which the Triton SDPA kernel does not take, so the CUDA
+    backend leaves that SDPA to the regular lowering.
     """
     print("\n" + "=" * 60)
     print("Exporting codec_decoder.pte")
@@ -155,13 +156,7 @@ def _export_codec_pte(model, args, device: str) -> None:
         qlinear_codec_group_size=args.qlinear_codec_group_size,
         device=device,
     )
-    codec_triton_mode = "OFF" if args.backend in ("cuda", "cuda-windows") else "ON"
-    et_codec = lower_to_executorch(
-        codec_programs,
-        codec_metadata,
-        backend=args.backend,
-        triton_kernel_mode=codec_triton_mode,
-    )
+    et_codec = lower_to_executorch(codec_programs, codec_metadata, backend=args.backend)
 
     codec_pte = os.path.join(args.output_dir, "codec_decoder.pte")
     print(f"\nSaving to {codec_pte}...")
@@ -475,15 +470,8 @@ def apply_model_quantization(
         )
 
 
-def lower_to_executorch(programs, metadata, backend="xnnpack", triton_kernel_mode="ON"):
-    """Lower exported programs to ExecuTorch.
-
-    Args:
-        triton_kernel_mode: For CUDA backend only. "ON" replaces ATen SDPA with
-            Triton sdpa kernel (required for the LM decoder). "OFF" disables
-            replacement so the codec's additive ALiBi mask SDPA can lower
-            (Triton sdpa kernel only accepts bool masks).
-    """
+def lower_to_executorch(programs, metadata, backend="xnnpack"):
+    """Lower exported programs to ExecuTorch."""
     mutable_buffer_passes = [InitializedMutableBufferPass(["k_cache", "v_cache"])]
     if backend == "xnnpack":
         from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
@@ -504,7 +492,7 @@ def lower_to_executorch(programs, metadata, backend="xnnpack", triton_kernel_mod
         print(
             f"\nLowering to ExecuTorch with CUDA "
             f"{'(Windows) ' if backend == 'cuda-windows' else ''}"
-            f"({len(programs)} methods, triton_kernel_mode={triton_kernel_mode})..."
+            f"({len(programs)} methods)..."
         )
         # NB: conv1d_to_conv2d is applied inside CudaBackend.preprocess via its
         # decomposition_table. Doing it here too triggers an extra run_decompositions
@@ -513,9 +501,6 @@ def lower_to_executorch(programs, metadata, backend="xnnpack", triton_kernel_mod
         partitioner = {}
         for key in programs:
             compile_specs = [CudaBackend.generate_method_name_compile_spec(key)]
-            compile_specs.append(
-                CompileSpec("triton_kernel_mode", triton_kernel_mode.encode("utf-8"))
-            )
             if backend == "cuda-windows":
                 compile_specs.append(CompileSpec("platform", b"windows"))
             partitioner[key] = [CudaPartitioner(compile_specs)]
@@ -540,7 +525,8 @@ def lower_to_executorch(programs, metadata, backend="xnnpack", triton_kernel_mod
             passes=mutable_buffer_passes,
             memory_planning_pass=MemoryPlanningPass(
                 alloc_graph_input=False,
-                share_mutable_buffers=True,
+                # The CUDA backend keeps mutable state inside its delegates.
+                share_mutable_buffers=backend not in ("cuda", "cuda-windows"),
             ),
             emit_mutable_buffer_names=True,
         ),