CUDA backend: lower rotary embeddings written with complex numbers

backends/aoti/aoti_backend.py

@@ -331,9 +331,9 @@ def preprocess(
                 custom_pass(device_edge_program.graph_module)
 
         # ``run_decompositions`` retraces the complete ExportedProgram even
-        # when none of the table's operators occur in the graph. The in-place
-        # passes above preserve graph inputs and outputs, so the existing graph
-        # signature remains valid when the decomposition table cannot apply.
+        # when none of the table's operators occur in the graph. Nothing else
+        # rebuilds the graph signature, so the in-place passes above must keep
+        # it valid when the decomposition table cannot apply.
         if contains_any_call_fn_target_op(
             device_edge_program.graph_module, decomposition_table
         ):

backends/cuda/BUCK

@@ -80,6 +80,7 @@ fbcode_target(
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
         "passes/remove_assert_async.py",
+        "passes/replace_complex_with_real.py",
         "passes/replace_fused_sdpa.py",
         "passes/replace_int64_floordiv.py",
         "passes/replace_masked_copy_with_where.py",

backends/cuda/cuda_backend.py

@@ -44,6 +44,9 @@
     MoveCondPredicateToCpuPass,
 )
 from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass
+from executorch.backends.cuda.passes.replace_complex_with_real import (
+    ReplaceComplexWithRealPass,
+)
 from executorch.backends.cuda.passes.replace_fused_sdpa import (
     ReplaceFusedSdpaWithSdpaPass,
 )
@@ -1266,6 +1269,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                 triton_kernel_mode = mode
         passes = [
             RemoveAssertAsyncPass(),
+            ReplaceComplexWithRealPass(),
             MoveCondPredicateToCpuPass(),
             ReplaceInt64FloorDivWithFloatPass(),
             ReplaceFusedSdpaWithSdpaPass(),