diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 64f0866a642..7240d057d70 100644 --- a/backends/cuda/cuda_backend.py +++ b/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,15 +1254,12 @@ 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() @@ -1269,7 +1267,16 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] 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 diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index f403dfb60ad..f3a889eb39e 100644 --- a/backends/cuda/tests/targets.bzl +++ b/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", diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 5861112dacc..7ec08c5f601 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/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 diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 176acd089e5..d20852d4bd3 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/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.""" diff --git a/examples/models/nemotron3_diarization/export_nemotron.py b/examples/models/nemotron3_diarization/export_nemotron.py index 4dba958add9..9197c81d93e 100644 --- a/examples/models/nemotron3_diarization/export_nemotron.py +++ b/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") } diff --git a/examples/models/voxtral_tts/export_voxtral_tts.py b/examples/models/voxtral_tts/export_voxtral_tts.py index 7ebc89bace1..fd5ec507975 100644 --- a/examples/models/voxtral_tts/export_voxtral_tts.py +++ b/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, ),