Add custom_op executor to the default executors to avoid hidden registration through thunder.torch.custom_op._register_custom_op (#2714)

Co-authored-by: Masaki Kozuki <mkozuki@nvidia.com>

Author: Copilot

thunder/__init__.py

@@ -199,9 +199,11 @@
 apex_executor: None | extend.Executor = extend.get_executor("apex")
 nvfuser_executor: None | extend.Executor = extend.get_executor("nvfuser")
 pytorch_executor: None | extend.Executor = extend.get_executor("torch")
+custom_op_executor: None | extend.Executor = extend.get_executor("custom_op")
 
-# Default executor list is [cudnn -> sdpa -> torchcompile_cat -> torchcompile_xentropy -> nvfuser -> torch -> python]
+# Default executor list is [cudnn -> sdpa -> torchcompile_cat -> torchcompile_xentropy -> nvfuser -> custom_op -> torch -> python]
 # Note that add_default_executor inserts executor at start of list, hence the reverse order below.
+add_default_executor(custom_op_executor)
 if nvfuser_executor:
     add_default_executor(nvfuser_executor)
 

thunder/extend/__init__.py

@@ -533,6 +533,7 @@ def get_all_executors() -> tuple[Executor, ...]:
         apexex,
         cudnn_layernormex,
         cudnnex,
+        custom_op_ex,
         nvfuserex,
         pythonex,
         sdpaex,

thunder/tests/test_recipes.py

@@ -15,7 +15,11 @@
 
 
 def get_expected_executors():
-    return [ex for ex in thunder.get_default_executors() if ex.name not in {"cudnn", "sdpa", "torchcompile_xentropy"}]
+    return [
+        ex
+        for ex in thunder.get_default_executors()
+        if ex.name not in {"cudnn", "sdpa", "torchcompile_xentropy", "custom_op"}
+    ]
 
 
 @pytest.mark.skipif(not nvfuser_available(), reason="nvFuser is not available")

thunder/torch/custom_op.py

@@ -332,9 +332,6 @@ def _register_custom_op(custom_op: CustomOpDef) -> Symbol:
     .. note::
         This feature is experimental and subject to change.
     """
-    from thunder.extend import add_executor_lists
-    from thunder.extend import get_default_executors
-    from thunder.extend import set_default_executors
     from thunder.executors.torchex import _always_executable
     from thunder.executors.custom_op_ex import custom_op_ex
     from thunder.torch import register_function
@@ -412,12 +409,6 @@ def _register_custom_op(custom_op: CustomOpDef) -> Symbol:
         backward_op = custom_op_ex.register_operator(bwd_fn_name, meta=backward_meta, fn=backward_impl)
         register_backward(symbol.id)(backward_op)
 
-    # NOTE: `thunder.extend.add_default_executor` basically does `lst.insert(ex, 0)`.
-    if custom_op_ex not in get_default_executors():
-        default_executors = get_default_executors()
-        new_default_executors = add_executor_lists(default_executors, [custom_op_ex])
-        set_default_executors(new_default_executors)
-
     _CUSTOM_OP_TO_TORCHFN_AND_SYMBOL[custom_op] = ((torch_opoverload, torch_opoverload_packet), symbol)
 
     return symbol