From dc30df3570e8fc2e8664082f7561921fc6f85147 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:51:00 -0700
Subject: [PATCH 1/5] Delete each CUDA output handle once when AOTI repeats it

When two outputs of a program are the same tensor, for example two views
of one buffer, AOTI puts the same handle in both output slots:

    output_handles[0] = buf0.release();
    output_handles[1] = output_handles[0];

The CUDA backend copied each output into its ExecuTorch tensor and then
deleted its handle, so the second copy read a freed handle and the
handle was deleted twice. The process crashed:

    SlimTensor storage_offset must not be negative, got -9007650542493611503
    Failed to D2D copy GPU output 1 into ETensor
    free(): double free detected in tcache 2

PyTorch's own AOTI runner handles this case. Now the CUDA backend gives
each repeated slot its own handle on the shared storage right after the
run, the same way aoti_torch_new_tensor_handle does, so no handle is
deleted twice.

The comment in ReplaceViewCopyWithViewPass no longer says that alias_copy
keeps two outputs on separate handles: two detached views of one tensor
become two alias_copy outputs, and they share one handle.

The new test returns two detached views of one tensor from a CUDA
program and compares both outputs with eager. It crashes without this
change.
---
 .../passes/replace_view_copy_with_view.py     |  4 +--
 backends/cuda/runtime/cuda_backend.cpp        | 12 +++++++++
 backends/cuda/tests/test_cuda_export.py       | 26 +++++++++++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/backends/aoti/passes/replace_view_copy_with_view.py b/backends/aoti/passes/replace_view_copy_with_view.py
index c881e650e7a..0c254543257 100644
--- a/backends/aoti/passes/replace_view_copy_with_view.py
+++ b/backends/aoti/passes/replace_view_copy_with_view.py
@@ -33,8 +33,8 @@
     ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor,
     # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either.
     # Without autograd it is ``alias_copy``, which Inductor can decompose,
-    # unlike ``detach_copy``. A view would let two outputs share one handle,
-    # and ``clone`` fails Edge verification in dim order mode.
+    # unlike ``detach_copy``, and ``clone`` fails Edge verification in dim
+    # order mode.
     torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default,
     ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default,
 }
diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp
index ce429a9a65a..88f8b1e3370 100644
--- a/backends/cuda/runtime/cuda_backend.cpp
+++ b/backends/cuda/runtime/cuda_backend.cpp
@@ -1025,6 +1025,18 @@ class ET_EXPERIMENTAL CudaBackend final
             nullptr);
     run_called = true;
 
+    // AOTI puts the same handle in every slot of an output that it returns
+    // more than once, such as two views of one buffer. Give each slot its own
+    // handle on the shared storage, so that each handle is deleted once.
+    for (size_t i = 0; i < n_outputs; i++) {
+      for (size_t j = 0; j < i; j++) {
+        if (slim_outputs[i] == slim_outputs[j]) {
+          slim_outputs[i] = new SlimTensor(*slim_outputs[j]);
+          break;
+        }
+      }
+    }
+
     // Delete orphaned pre-created outputs that run() replaced.
     // Must happen before the error check — if run() fails after
     // replacing some outputs, the originals would otherwise leak.
diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py
index 637ff4036b1..5ca9b52fdb8 100644
--- a/backends/cuda/tests/test_cuda_export.py
+++ b/backends/cuda/tests/test_cuda_export.py
@@ -766,6 +766,32 @@ def forward(self, x, mask, column, idx, counts):
                 for output, expected in zip(outputs, model(*inputs), strict=True):
                     torch.testing.assert_close(output.cpu(), expected)
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_same_tensor_returned_twice(self):
+        """AOTI returns one handle for the two outputs, since both are views
+        of the same buffer."""
+        from executorch.runtime import Runtime
+
+        class TwoViews(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                y = x * 2
+                return y.detach(), y.detach()
+
+        model = TwoViews()
+        x = torch.randn(4, 8)
+        edge = self._export_to_cuda_with_lower(model, (x,))
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.pte")
+            with open(path, "wb") as f:
+                edge.to_executorch().write_to_file(f)
+            method = Runtime.get().load_program(path).load_method("forward")
+            outputs = method.execute([x])
+        for output, expected in zip(outputs, model(x), strict=True):
+            torch.testing.assert_close(output.cpu(), expected)
+
     def test_device_info_propagated_to_cuda_delegate_outputs(self):
         """
         Verify that, for a CUDA-delegated graph, every memory-planned tensor's

From 9181a61aaa7672879ebac422e49237948614ac97 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:33:47 -0700
Subject: [PATCH 2/5] Lower rotary embeddings written with complex numbers in
 the CUDA backend

Rotary position embeddings are often written with complex numbers:
torch.polar builds the rotations, torch.view_as_complex turns pairs of
values into complex numbers, they are multiplied, and torch.view_as_real
turns the result back. Such a model exports, but the CUDA backend could
not lower it:

    RuntimeError: Method forward missing fallback kernels (5 total):
      - aoti_torch_cuda_mul_Tensor
      - aoti_torch_cuda_polar
      - aten::view_as_complex_copy
      - aten::view_as_real_copy
      - prims::convert_element_type

Inductor generates no code for complex numbers. It calls an ATen kernel
for every op that reads or writes a complex tensor, and this backend
ships none of those kernels. Torch-TensorRT leaves polar and view_as_real
to the CUDA backend, so its programs stop at the same error.

A new pass, ReplaceComplexWithRealPass, keeps each complex tensor as a
real tensor with a last dim of size 2, the layout of view_as_real, and
rewrites the ops on it as real math before Inductor compiles the graph.
polar becomes cosine and sine times the magnitude, a complex product
becomes four real products, and view_as_complex and view_as_real go
away. Views, slices, splits, indexing, repeats, transposes, unbinding,
concatenation, stacking and casts keep their op, with the dims and dtype
of the pair. A cast from real to complex gets zero imaginary parts, and
a complex constant, buffer or parameter is stored as its real view.

A decomposition table cannot do this, because a decomposition returns
the dtype of the op it replaces.

The pass rewrites a program only when it can rewrite every op that
reads or writes a complex tensor in it. Otherwise, for example for a
complex add, a complex program input, complex32 or complex values inside
a torch.cond branch, it leaves the program as it is and logs the first op
it cannot rewrite.

Complex128 math becomes float64 math, so a program that keeps float64
values in memory also needs float64 support in the runtime.

test_complex_rotary_embedding lowers a rotary embedding with a complex
constant, polar, views, a split, a concatenation and a view_as_real
output, runs it with the runtime and compares it with eager, in float32
and bfloat16. On the CPU, with the program prepared the way the CUDA
backend's passes get it, test_complex_rewrite checks a rotary embedding
with a complex constant and parameter, a gather by position, views,
splits and joins of many kinds and a dynamic length against eager, after
the retrace that the backend runs when its decomposition table applies,
and test_complex_rewrite_leaves_unsupported_programs checks that the
programs listed above, a channels-last copy, a cast to real and a
mutated complex buffer stay as they are. They fail without this change.
---
 backends/aoti/aoti_backend.py                 |   4 +-
 backends/cuda/BUCK                            |   1 +
 backends/cuda/cuda_backend.py                 |   9 +-
 .../cuda/passes/replace_complex_with_real.py  | 384 ++++++++++++++++++
 backends/cuda/tests/test_cuda_export.py       | 198 ++++++++-
 5 files changed, 592 insertions(+), 4 deletions(-)
 create mode 100644 backends/cuda/passes/replace_complex_with_real.py

diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py
index 0d0670b7318..bb6aa12e29a 100644
--- a/backends/aoti/aoti_backend.py
+++ b/backends/aoti/aoti_backend.py
@@ -312,8 +312,8 @@ def preprocess(
 
         # ``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.
+        # passes above keep the graph signature valid, so the program needs no
+        # retrace when the decomposition table cannot apply.
         if contains_any_call_fn_target_op(
             device_edge_program.graph_module, decomposition_table
         ):
diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK
index ca8f9eee753..b2d0ba15a5d 100644
--- a/backends/cuda/BUCK
+++ b/backends/cuda/BUCK
@@ -79,6 +79,7 @@ fbcode_target(
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
+        "passes/replace_complex_with_real.py",
         "passes/replace_int64_floordiv.py",
     ],
     visibility = [
diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py
index 891fb81bd06..db0b222822d 100644
--- a/backends/cuda/cuda_backend.py
+++ b/backends/cuda/cuda_backend.py
@@ -39,6 +39,9 @@
 from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
     MoveCondPredicateToCpuPass,
 )
+from executorch.backends.cuda.passes.replace_complex_with_real import (
+    ReplaceComplexWithRealPass,
+)
 from executorch.backends.cuda.passes.replace_int64_floordiv import (
     ReplaceInt64FloorDivWithFloatPass,
 )
@@ -1013,7 +1016,11 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                         f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 triton_kernel_mode = mode
-        passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        passes = [
+            ReplaceComplexWithRealPass(),
+            MoveCondPredicateToCpuPass(),
+            ReplaceInt64FloorDivWithFloatPass(),
+        ]
         cell_layout = any(
             spec.key == OFFGRAPH_KV_COMPILE_SPEC
             and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"
diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py
new file mode 100644
index 00000000000..00cf34d4717
--- /dev/null
+++ b/backends/cuda/passes/replace_complex_with_real.py
@@ -0,0 +1,384 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Graph Transformation Pass for Complex Numbers.
+
+Inductor generates no code for complex numbers. It calls an ATen kernel for
+every op that reads or writes a complex tensor, and this backend ships none of
+those kernels. So this pass keeps each complex tensor as a real tensor with a
+last dim of size 2 that holds the real and imaginary parts, the layout of
+``torch.view_as_real``, and rewrites the ops on it as ops on real numbers:
+
+    view_as_complex(x)  ->  x
+    view_as_real(z)     ->  z
+    polar(r, theta)     ->  stack([r * cos(theta), r * sin(theta)], -1)
+    z * w               ->  stack([zr * wr - zi * wi, zr * wi + zi * wr], -1)
+
+Views, copies, indexing, repeats, splits, concatenation and stacking act on
+both parts alike, so they keep their op, with the dims they name counted in
+the complex tensor. A complex constant is stored as its real view. A
+decomposition table cannot do this, since a decomposition returns the dtype of
+the op it replaces.
+
+The pass rewrites a program only when it can rewrite every op that reads or
+writes a complex tensor in it, and it does not rewrite control flow branches,
+such as those of torch.cond. Otherwise it leaves the program as it is and logs
+the first op it cannot rewrite, and the backend reports the ops it has no
+kernels for.
+"""
+
+import logging
+import operator
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
+from torch.fx import Graph, Node
+
+logger = logging.getLogger(__name__)
+
+_aten = exir_ops.edge.aten
+
+
+def _complex_vals(node: Node) -> List[torch.Tensor]:
+    val = node.meta.get("val")
+    vals = val if isinstance(val, (list, tuple)) else (val,)
+    return [v for v in vals if isinstance(v, torch.Tensor) and v.is_complex()]
+
+
+# The pair's extra dim comes last, so a dim of the complex tensor names the
+# same dim of the pair once it is counted from the front.
+def _dim(dim: int, ndim: int) -> int:
+    return dim + ndim if dim < 0 else dim
+
+
+def _dims(dims: List[int], ndim: int) -> List[int]:
+    return [_dim(d, ndim) for d in dims]
+
+
+def _new_dim(dim: int, ndim: int) -> int:
+    return _dim(dim, ndim + 1)
+
+
+def _permutation(dims: List[int], ndim: int) -> List[int]:
+    return _dims(dims, ndim) + [ndim]
+
+
+def _repeats(repeats: List[Any], ndim: int) -> List[Any]:
+    return [*repeats, 1]
+
+
+def _size(size: List[Any], ndim: int) -> List[Any]:
+    return [*size, 2]
+
+
+# Ops that act on the real and imaginary parts alike. For each argument that
+# names dims or sizes of the complex tensor, the entry gives the argument's
+# position, and how to name the same dims of the pair. These arguments'
+# default values never need changing.
+_SAME_ON_BOTH_PARTS: Dict[Any, Tuple[Tuple[int, Callable[..., Any]], ...]] = {
+    _aten.alias_copy.default: (),
+    _aten.cat.default: ((1, _dim),),
+    _aten.clone.default: (),
+    _aten.expand_copy.default: ((1, _size),),
+    _aten.index.Tensor: (),
+    _aten.permute_copy.default: ((1, _permutation),),
+    _aten.repeat.default: ((1, _repeats),),
+    _aten.select.int: ((1, _dim),),
+    _aten.select_copy.int: ((1, _dim),),
+    _aten.slice.Tensor: ((1, _dim),),
+    _aten.slice_copy.Tensor: ((1, _dim),),
+    _aten.split.Tensor: ((2, _dim),),
+    _aten.split_copy.Tensor: ((2, _dim),),
+    _aten.split_with_sizes_copy.default: ((2, _dim),),
+    _aten.squeeze_copy.default: (),
+    _aten.squeeze_copy.dim: ((1, _dim),),
+    _aten.squeeze_copy.dims: ((1, _dims),),
+    _aten.stack.default: ((1, _new_dim),),
+    _aten.transpose_copy.int: ((1, _dim), (2, _dim)),
+    _aten.unbind_copy.int: ((1, _dim),),
+    _aten.unsqueeze_copy.default: ((1, _new_dim),),
+    _aten.view_copy.default: ((1, _size),),
+}
+
+_LIST_INPUTS = (_aten.cat.default, _aten.stack.default)
+
+# Copies that may change the dtype. The pair gets the precision of the
+# complex dtype, and the complex tensor's dim order with its extra dim last.
+_COPIES = (
+    _aten._to_copy.default,
+    exir_ops.edge.dim_order_ops._clone_dim_order.default,
+    exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+)
+_DENSE_FORMATS = (None, torch.preserve_format, torch.contiguous_format)
+_PROVENANCE = ("from_node", "nn_module_stack", "source_fn_stack", "stack_trace")
+
+# The fake kernels of some view copies, such as permute_copy, make dynamic
+# sizes static or fail on them, so a new view copy gets its value from its
+# view. view_copy keeps its own kernel, which keeps dynamic sizes, because
+# view fails on some strides that view_copy accepts.
+_VIEWS: Dict[Any, Any] = {
+    _aten.alias_copy.default: torch.ops.aten.alias.default,
+    _aten.expand_copy.default: torch.ops.aten.expand.default,
+    _aten.permute_copy.default: torch.ops.aten.permute.default,
+    _aten.select_copy.int: torch.ops.aten.select.int,
+    _aten.slice_copy.Tensor: torch.ops.aten.slice.Tensor,
+    _aten.split_copy.Tensor: torch.ops.aten.split.Tensor,
+    _aten.split_with_sizes_copy.default: torch.ops.aten.split_with_sizes.default,
+    _aten.squeeze_copy.default: torch.ops.aten.squeeze.default,
+    _aten.squeeze_copy.dim: torch.ops.aten.squeeze.dim,
+    _aten.squeeze_copy.dims: torch.ops.aten.squeeze.dims,
+    _aten.transpose_copy.int: torch.ops.aten.transpose.int,
+    _aten.unbind_copy.int: torch.ops.aten.unbind.int,
+    _aten.unsqueeze_copy.default: torch.ops.aten.unsqueeze.default,
+}
+
+
+class _Rewriter:
+    """Builds, in graph order, the real replacement of every node that reads
+    or writes a complex tensor."""
+
+    def __init__(self, graph: Graph, complex_nodes: Set[Node]) -> None:
+        self.graph = graph
+        self.complex_nodes = complex_nodes
+        self.ndims = {
+            node: node.meta["val"].dim()
+            for node in complex_nodes
+            if isinstance(node.meta["val"], torch.Tensor)
+        }
+        self.replacements: Dict[Node, Node] = {}
+        self.source: Optional[Node] = None
+
+    def get(self, value: Any) -> Any:
+        """``value`` with each complex node in it replaced by its pair."""
+        return torch.fx.map_arg(
+            value,
+            lambda n: self.replacements[n] if n in self.complex_nodes else n,
+        )
+
+    def call(self, target: Any, *args: Any, **kwargs: Any) -> Node:
+        node = self.graph.call_function(target, args, kwargs)
+        # Errors and profiles then point at the op that this one replaces.
+        node.meta.update(
+            {k: v for k, v in self.source.meta.items() if k in _PROVENANCE}
+        )
+        fake_args, fake_kwargs = torch.fx.map_arg(
+            (args, kwargs), lambda n: n.meta["val"]
+        )
+        node.meta["val"] = _VIEWS.get(target, target)(*fake_args, **fake_kwargs)
+        return node
+
+    def parts(self, value: Node) -> List[Node]:
+        pair = self.get(value)
+        return [self.call(_aten.select.int, pair, -1, i) for i in (0, 1)]
+
+
+def _view_as_complex(rw: _Rewriter, node: Node) -> Node:
+    return node.args[0]
+
+
+def _view_as_real(rw: _Rewriter, node: Node) -> Node:
+    return rw.get(node.args[0])
+
+
+def _polar(rw: _Rewriter, node: Node) -> Node:
+    magnitude, angle = node.args
+    parts = [
+        rw.call(_aten.mul.Tensor, magnitude, rw.call(trig, angle))
+        for trig in (_aten.cos.default, _aten.sin.default)
+    ]
+    return rw.call(_aten.stack.default, parts, -1)
+
+
+def _mul(rw: _Rewriter, node: Node) -> Node:
+    (ar, ai), (br, bi) = (rw.parts(arg) for arg in node.args)
+    mul = _aten.mul.Tensor
+    real = rw.call(_aten.sub.Tensor, rw.call(mul, ar, br), rw.call(mul, ai, bi))
+    imag = rw.call(_aten.add.Tensor, rw.call(mul, ar, bi), rw.call(mul, ai, br))
+    return rw.call(_aten.stack.default, [real, imag], -1)
+
+
+def _copy(rw: _Rewriter, node: Node) -> Node:
+    source, kwargs = node.args[0], dict(node.kwargs)
+    if kwargs.get("dtype") is not None:
+        kwargs["dtype"] = kwargs["dtype"].to_real()
+    if source not in rw.complex_nodes:
+        # Edge programs cast a real operand to complex before complex math
+        # with it. The cast gives zero imaginary parts.
+        real = rw.call(node.target, source, **kwargs)
+        imag = rw.call(_aten.zeros_like.default, real)
+        return rw.call(_aten.stack.default, [real, imag], -1)
+    if kwargs.get("dim_order") is not None:
+        kwargs["dim_order"] = [*kwargs["dim_order"], len(kwargs["dim_order"])]
+    return rw.call(node.target, rw.get(source), **kwargs)
+
+
+def _same_on_both_parts(rw: _Rewriter, node: Node) -> Node:
+    args = list(node.args)
+    tensor = args[0][0] if node.target in _LIST_INPUTS else args[0]
+    for index, pair_value in _SAME_ON_BOTH_PARTS[node.target]:
+        if index < len(args):
+            args[index] = pair_value(args[index], rw.ndims[tensor])
+    return rw.call(node.target, *rw.get(args), **rw.get(dict(node.kwargs)))
+
+
+def _getitem(rw: _Rewriter, node: Node) -> Node:
+    return rw.call(operator.getitem, rw.get(node.args[0]), node.args[1])
+
+
+_REWRITES: Dict[Any, Callable[[_Rewriter, Node], Node]] = {
+    _aten.view_as_complex_copy.default: _view_as_complex,
+    _aten.view_as_real_copy.default: _view_as_real,
+    _aten.polar.default: _polar,
+    _aten.mul.Tensor: _mul,
+    operator.getitem: _getitem,
+    **{op: _copy for op in _COPIES},
+    **{op: _same_on_both_parts for op in _SAME_ON_BOTH_PARTS},
+}
+
+
+def _keeps_dims_in_order(node: Node) -> bool:
+    # With the pair's dim after them, the dims of a channels-last layout are
+    # in an order that no memory format describes, and the dim order ops
+    # accept only those that one does.
+    dim_order = node.kwargs.get("dim_order")
+    return node.kwargs.get("memory_format") in _DENSE_FORMATS and (
+        dim_order is None or list(dim_order) == list(range(len(dim_order)))
+    )
+
+
+def _can_rewrite(node: Node, complex_nodes: Set[Node]) -> bool:
+    if node.op != "call_function" or node.target not in _REWRITES:
+        return False
+    # The rewrites read by position the arguments that exported programs
+    # pass by position.
+    schema = getattr(node.target, "_schema", None)
+    if schema is not None and any(
+        not arg.kwarg_only for arg in schema.arguments if arg.name in node.kwargs
+    ):
+        return False
+    if node.target in _COPIES:
+        # A cast from complex to real has no rewrite here.
+        return node in complex_nodes and _keeps_dims_in_order(node)
+    if node.target is _aten.clone.default:
+        return _keeps_dims_in_order(node)
+    if node.target is _aten.mul.Tensor:
+        # Edge programs cast a real factor to complex first.
+        return all(arg in complex_nodes for arg in node.args)
+    if node.target in _LIST_INPUTS:
+        return all(tensor in complex_nodes for tensor in node.args[0])
+    return True
+
+
+def _constant_name(exported_program: ExportedProgram, node: Node) -> Optional[str]:
+    """The name of the constant, parameter or buffer that ``node`` reads, if
+    the graph never writes it."""
+    signature = exported_program.graph_signature
+    for inputs in (
+        signature.inputs_to_lifted_tensor_constants,
+        signature.inputs_to_parameters,
+    ):
+        if node.name in inputs:
+            return inputs[node.name]
+    buffer = signature.inputs_to_buffers.get(node.name)
+    if buffer in signature.buffers_to_mutate.values():
+        return None
+    return buffer
+
+
+def _store_as_real(exported_program: ExportedProgram, node: Node) -> None:
+    name = _constant_name(exported_program, node)
+    tensors = (
+        exported_program._constants
+        if name in exported_program.constants
+        else exported_program._state_dict
+    )
+    tensor = tensors[name]
+    pair = torch.view_as_real(tensor.detach().resolve_conj()).clone()
+    if isinstance(tensor, torch.nn.Parameter):
+        pair = torch.nn.Parameter(pair, requires_grad=False)
+    tensors[name] = pair
+    node.meta["val"] = torch.view_as_real(node.meta["val"].resolve_conj())
+
+
+def _first_unsupported(
+    exported_program: ExportedProgram, complex_nodes: Set[Node], affected: Set[Node]
+) -> Optional[Node]:
+    """The first node that reads or writes a complex tensor and that the pass
+    cannot rewrite, if there is one."""
+    graph_module = exported_program.graph_module
+    for module in graph_module.modules():
+        if module is not graph_module and isinstance(module, torch.fx.GraphModule):
+            for node in module.graph.nodes:
+                if _complex_vals(node):
+                    return node
+    for node in graph_module.graph.nodes:
+        if node not in affected:
+            continue
+        if node.op == "placeholder":
+            supported = _constant_name(exported_program, node) is not None
+        else:
+            supported = _can_rewrite(node, complex_nodes)
+        # PyTorch computes complex32 math in float32, and its pair would
+        # compute in float16.
+        if not supported or any(
+            val.dtype == torch.complex32 for val in _complex_vals(node)
+        ):
+            return node
+    return None
+
+
+class ReplaceComplexWithRealPass:
+    """
+    Replace complex tensors with real tensors that hold their real and
+    imaginary parts in a last dim of size 2, and the ops on them with ops on
+    real numbers, so that Inductor can generate code for them. See the module
+    docstring for the ops it handles.
+    """
+
+    requires_exported_program = True
+
+    def __call__(self, exported_program: ExportedProgram) -> None:
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        complex_nodes = {node for node in graph.nodes if _complex_vals(node)}
+        affected = complex_nodes | {
+            user for node in complex_nodes for user in node.users
+        }
+        unsupported = _first_unsupported(exported_program, complex_nodes, affected)
+        if unsupported is not None:
+            logger.warning(
+                "The CUDA backend cannot lower the complex numbers in this "
+                "program. The first op it cannot rewrite: %s",
+                unsupported.format_node(),
+            )
+            return
+        if not complex_nodes:
+            return
+
+        rw = _Rewriter(graph, complex_nodes)
+        order = [node for node in graph.nodes if node in affected]
+        for node in order:
+            if node.op == "placeholder":
+                _store_as_real(exported_program, node)
+                rw.replacements[node] = node
+                continue
+            rw.source = node
+            with graph.inserting_before(node):
+                rw.replacements[node] = _REWRITES[node.target](rw, node)
+        signature = exported_program.graph_signature
+        for node in reversed(order):
+            replacement = rw.replacements[node]
+            if replacement is not node:
+                node.replace_all_uses_with(replacement)
+                # The signature names the program's outputs by node, and a
+                # view_as_real can be one of them.
+                signature.replace_all_uses(node.name, replacement.name)
+                graph.erase_node(node)
+        graph.lint()
+        graph_module.recompile()
diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py
index 5ca9b52fdb8..7298b6bc037 100644
--- a/backends/cuda/tests/test_cuda_export.py
+++ b/backends/cuda/tests/test_cuda_export.py
@@ -17,11 +17,36 @@
 from executorch.backends.cuda.cuda_backend import CudaBackend
 from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
 from executorch.examples.models.toy_model import SdpaModule
-from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower
+from executorch.exir import (
+    EdgeCompileConfig,
+    schema,
+    to_edge,
+    to_edge_transform_and_lower,
+)
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export import export
 
 
+def _program_for_cuda_passes(model, inputs, dynamic_shapes=None):
+    """The Edge program of ``model`` as the CUDA backend's passes get it: no
+    op decomposed, as the CUDA partitioner asks, and views for copies."""
+    from executorch.backends.aoti.passes.replace_view_copy_with_view import (
+        ReplaceViewCopyWithViewPass,
+    )
+
+    program = export(model, inputs, dynamic_shapes=dynamic_shapes, strict=True)
+    # to_edge_transform_and_lower functionalizes the program before it asks
+    # the partitioner which ops to keep.
+    program = program.run_decompositions({})
+    ops, _ = CudaPartitioner([]).ops_to_not_decompose(program)
+    program = to_edge(
+        program,
+        compile_config=EdgeCompileConfig(_check_ir_validity=False, preserve_ops=ops),
+    ).exported_program()
+    ReplaceViewCopyWithViewPass()(program.graph_module)
+    return program
+
+
 class TestCudaBackendCompileOptions(unittest.TestCase):
     def test_low_memory_triton_reduction_loads_stay_loop_scoped(self):
         from executorch.backends.cuda.cuda_backend import (
@@ -291,6 +316,117 @@ def specs(name):
                 {"forward": specs("first"), "second": specs("second")},
             )
 
+    def test_complex_rewrite(self):
+        """The rewrite of complex numbers gives eager's results."""
+        from executorch.backends.cuda.passes.replace_complex_with_real import (
+            ReplaceComplexWithRealPass,
+        )
+
+        angles = torch.outer(torch.arange(8.0), torch.rand(3))
+
+        class Rotate(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                # A plain attribute becomes a constant of the program.
+                self.freqs = torch.polar(torch.ones_like(angles), angles)
+                self.phase = torch.nn.Parameter(
+                    torch.polar(torch.ones(2, 6), torch.rand(2, 6)), requires_grad=False
+                )
+
+            def forward(self, x, positions, scale):
+                f = torch.stack(self.freqs[positions].unbind(-1), -1)
+                f = f.unsqueeze(-2).expand(-1, 2, -1).flatten(-2).unsqueeze(-2)
+                pairs = torch.view_as_complex(x.unflatten(-1, (-1, 2)))
+                # The product keeps the permuted layout, so export flattens
+                # the transpose back with a view.
+                scaled = pairs.unflatten(-1, (2, 3)).permute(0, 1, -1, -2) * scale
+                scaled = scaled.transpose(-1, -2).flatten(-2)
+                phase = self.phase.select(-2, 1).repeat(1, 1).detach()
+                rotated = scaled * f.to(torch.complex128) * phase
+                first, second = rotated.chunk(2, dim=-1)
+                second = torch.cat(second.split([1, 2], dim=-1), -1)
+                return torch.view_as_real(
+                    torch.cat([first, second], -1).squeeze(-2).clone()
+                )
+
+        model = Rotate()
+        scale = torch.randn(3, 2)
+        x = torch.randn(5, 1, 12, dtype=torch.float64)
+        length = torch.export.Dim("length", min=2, max=64)
+        program = _program_for_cuda_passes(
+            model,
+            (x, torch.tensor([3, 0, 7, 7, 1]), scale),
+            dynamic_shapes=({0: length}, {0: length}, None),
+        )
+        ReplaceComplexWithRealPass()(program)
+        for node in program.graph.nodes:
+            val = node.meta.get("val")
+            self.assertFalse(isinstance(val, torch.Tensor) and val.is_complex())
+        # The backend retraces the program when its decomposition table applies.
+        program = program.run_decompositions({})
+        for positions in (torch.tensor([3, 0, 7, 7, 1]), torch.tensor([6, 2, 4])):
+            inputs = (x[: len(positions)], positions, scale)
+            torch.testing.assert_close(program.module()(*inputs), model(*inputs))
+        self.assertTrue(
+            any(
+                isinstance(p, ReplaceComplexWithRealPass)
+                for p in CudaBackend.get_custom_passes([])
+            )
+        )
+
+    def test_complex_rewrite_leaves_unsupported_programs(self):
+        """A program with a complex op that the rewrite does not handle keeps
+        all its complex numbers, so no program is left half rewritten."""
+        from executorch.backends.cuda.passes.replace_complex_with_real import (
+            ReplaceComplexWithRealPass,
+        )
+
+        class MutatedBuffer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("freqs", torch.randn(3, dtype=torch.cfloat))
+
+            def forward(self, x):
+                self.freqs.mul_(torch.view_as_complex(x))
+                return torch.view_as_real(self.freqs.clone())
+
+        class Fn(torch.nn.Module):
+            def __init__(self, fn):
+                super().__init__()
+                self.fn = fn
+
+            def forward(self, *args):
+                return self.fn(*args)
+
+        def square(x):
+            z = torch.view_as_complex(x)
+            return torch.view_as_real(z * z)
+
+        def channels_last(x):
+            z = torch.view_as_complex(x).clone(memory_format=torch.channels_last)
+            return torch.view_as_real(z)
+
+        def in_branch(x):
+            return torch.cond(x.sum() > 0, square, torch.neg, (square(x),))
+
+        pairs = torch.randn(1, 2, 3, 4, 2)
+        for model, inputs in (
+            (Fn(lambda z: torch.view_as_real(z * z)), (torch.randn(3).cfloat(),)),
+            (Fn(lambda x: torch.view_as_real(torch.view_as_complex(x) + 1)), (pairs,)),
+            (Fn(square), (pairs.half(),)),
+            (Fn(lambda x: torch.view_as_complex(x).float()), (pairs,)),
+            (Fn(channels_last), (pairs,)),
+            (Fn(in_branch), (pairs,)),
+            (MutatedBuffer(), (torch.randn(3, 2),)),
+        ):
+            program = _program_for_cuda_passes(model, inputs)
+            code = program.graph_module.code
+            tensors = {**program.state_dict, **program.constants}
+            ReplaceComplexWithRealPass()(program)
+            self.assertEqual(program.graph_module.code, code)
+            for name, tensor in {**program.state_dict, **program.constants}.items():
+                self.assertIs(tensor, tensors[name])
+
 
 class TestCudaExport(unittest.TestCase):
     """Test CUDA export functionality for various operations using to_edge_transform_and_lower."""
@@ -792,6 +928,66 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         for output, expected in zip(outputs, model(x), strict=True):
             torch.testing.assert_close(output.cpu(), expected)
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_complex_rotary_embedding(self):
+        """Inductor generates no code for complex numbers, so the backend
+        computes a rotary embedding written with them on real numbers."""
+        from executorch.runtime import Runtime
+
+        class ComplexRope(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                angles = torch.outer(torch.arange(4.0), torch.rand(12))
+                self.register_buffer(
+                    "freqs", torch.polar(torch.ones_like(angles), angles)
+                )
+
+            def forward(self, x, angles):
+                # A grid of positions takes each axis's part of the frequencies.
+                rows, cols = self.freqs.split([8, 4], dim=1)
+                grid = torch.cat(
+                    [
+                        rows[:2].view(2, 1, -1).expand(2, 2, -1),
+                        cols[:2].view(1, 2, -1).expand(2, 2, -1),
+                    ],
+                    dim=-1,
+                ).view(4, 1, -1)
+                spin = torch.polar(torch.ones_like(angles), angles).unsqueeze(1)
+                pairs = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+                rotated = torch.view_as_real(pairs * grid * spin * 0.5)
+                # A view_as_real can also be an output of the program.
+                return rotated.flatten(-2).type_as(x), torch.view_as_real(spin)
+
+        model = ComplexRope()
+        angles = torch.randn(4, 12)
+        for dtype in (torch.float32, torch.bfloat16):
+            with self.subTest(dtype=dtype):
+                x = torch.randn(4, 3, 24, dtype=dtype)
+                program = self._export_to_cuda_with_lower(
+                    model, (x, angles)
+                ).to_executorch()
+                with tempfile.TemporaryDirectory() as directory:
+                    path = os.path.join(directory, "model.pte")
+                    with open(path, "wb") as f:
+                        program.write_to_file(f)
+                    program.write_tensor_data_to_file(directory)
+                    (data,) = [
+                        os.path.join(directory, name)
+                        for name in os.listdir(directory)
+                        if name.endswith(".ptd")
+                    ]
+                    method = (
+                        Runtime.get()
+                        .load_program(path, data_path=data)
+                        .load_method("forward")
+                    )
+                    outputs = method.execute([x, angles])
+                for output, expected in zip(outputs, model(x, angles), strict=True):
+                    torch.testing.assert_close(output.cpu(), expected)
+
     def test_device_info_propagated_to_cuda_delegate_outputs(self):
         """
         Verify that, for a CUDA-delegated graph, every memory-planned tensor's

From 7592e129d82668dba67671598aca12af791b918d Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 12:06:38 -0800
Subject: [PATCH 3/5] Lower fused attention variants in the CUDA backend

Decomposing scaled_dot_product_attention on a CUDA device turns it into
one of its fused variants: _scaled_dot_product_flash_attention without a
mask, or _scaled_dot_product_efficient_attention with one. torch.export's
default decompositions do this, and so does Torch-TensorRT for an
attention it leaves to the CUDA backend, for example one whose keys are a
constant. The CUDA backend had no kernel for these ops:

    RuntimeError: Method forward missing fallback kernels (1 total):
      - at::_ops::_scaled_dot_product_flash_attention::call

A new pass, ReplaceFusedSdpaWithSdpaPass, turns the flash, efficient and
cuDNN variants back into scaled_dot_product_attention with the same
query, key, value, bias, dropout, causal flag and scale, and with
enable_gqa when there are fewer key heads than query heads. The Triton
SDPA kernel or Inductor's math decomposition then lowers it, as it does
any other attention. The pass also looks inside control flow branches.

The pass rewrites a node only when scaled_dot_product_attention computes
its one used output. Its other outputs, such as the log-sum-exp, must be
unused, which is the case for inference. A causal call must have no
bias, and a causal flash call must have equal query and key lengths,
since flash aligns the causal mask to the bottom right and
scaled_dot_product_attention to the top left. Otherwise the node stays,
and the error above still names it.

test_fused_sdpa lowers a program with a flash attention (fewer key
heads, causal, a custom scale) and an efficient attention with a bias,
runs it with the runtime and compares both outputs with eager. It also
checks that the pass leaves a causal flash attention with fewer queries
than keys. It fails without this change.
---
 backends/cuda/BUCK                         |   1 +
 backends/cuda/cuda_backend.py              |   4 +
 backends/cuda/passes/replace_fused_sdpa.py | 131 +++++++++++++++++++++
 backends/cuda/tests/test_cuda_export.py    |  52 ++++++++
 4 files changed, 188 insertions(+)
 create mode 100644 backends/cuda/passes/replace_fused_sdpa.py

diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK
index b2d0ba15a5d..a0e544de388 100644
--- a/backends/cuda/BUCK
+++ b/backends/cuda/BUCK
@@ -80,6 +80,7 @@ fbcode_target(
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
         "passes/replace_complex_with_real.py",
+        "passes/replace_fused_sdpa.py",
         "passes/replace_int64_floordiv.py",
     ],
     visibility = [
diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py
index db0b222822d..1ec28eb9b2c 100644
--- a/backends/cuda/cuda_backend.py
+++ b/backends/cuda/cuda_backend.py
@@ -42,6 +42,9 @@
 from executorch.backends.cuda.passes.replace_complex_with_real import (
     ReplaceComplexWithRealPass,
 )
+from executorch.backends.cuda.passes.replace_fused_sdpa import (
+    ReplaceFusedSdpaWithSdpaPass,
+)
 from executorch.backends.cuda.passes.replace_int64_floordiv import (
     ReplaceInt64FloorDivWithFloatPass,
 )
@@ -1020,6 +1023,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
             ReplaceComplexWithRealPass(),
             MoveCondPredicateToCpuPass(),
             ReplaceInt64FloorDivWithFloatPass(),
+            ReplaceFusedSdpaWithSdpaPass(),
         ]
         cell_layout = any(
             spec.key == OFFGRAPH_KV_COMPILE_SPEC
diff --git a/backends/cuda/passes/replace_fused_sdpa.py b/backends/cuda/passes/replace_fused_sdpa.py
new file mode 100644
index 00000000000..04b425eb4ec
--- /dev/null
+++ b/backends/cuda/passes/replace_fused_sdpa.py
@@ -0,0 +1,131 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Graph Transformation Pass for Fused Attention Replacement.
+
+Decomposing ``scaled_dot_product_attention`` on CUDA, as ``torch.export``'s
+default decompositions and Torch-TensorRT do, turns it into one of its fused
+CUDA variants. The backend has no kernel for those, so this pass turns them
+back into ``scaled_dot_product_attention``, which the Triton SDPA kernel or
+Inductor's math decomposition lowers.
+"""
+
+import logging
+import operator
+from typing import Any, Dict, Optional
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
+from torch.export.graph_signature import ExportGraphSignature
+from torch.fx import GraphModule, Node
+from torch.fx.experimental.symbolic_shapes import statically_known_true
+
+logger = logging.getLogger(__name__)
+
+_FLASH = exir_ops.edge.aten._scaled_dot_product_flash_attention.default
+_FUSED_SDPA_OPS = (
+    _FLASH,
+    exir_ops.edge.aten._scaled_dot_product_efficient_attention.default,
+    exir_ops.edge.aten._scaled_dot_product_cudnn_attention.default,
+)
+
+
+class ReplaceFusedSdpaWithSdpaPass:
+    """
+    Replace each fused SDPA node with ``scaled_dot_product_attention`` when
+    that computes its only used output. The other outputs, such as the
+    log-sum-exp, are only for the backward pass. Any other node stays as it
+    is.
+    """
+
+    requires_exported_program = True
+
+    def __call__(self, exported_program: ExportedProgram) -> None:
+        program_module = exported_program.graph_module
+        for module in program_module.modules():
+            if not isinstance(module, GraphModule):
+                continue
+            # The signature names the outputs of the program's own graph only.
+            signature = (
+                exported_program.graph_signature if module is program_module else None
+            )
+            modified = False
+            for node in list(module.graph.nodes):
+                if node.op != "call_function" or node.target not in _FUSED_SDPA_OPS:
+                    continue
+                values = _arguments(node)
+                if _sdpa_computes(node, values):
+                    _replace_with_sdpa(module, node, values, signature)
+                    modified = True
+                else:
+                    logger.info(f"Keeping {node.name}: SDPA computes it differently")
+            if modified:
+                module.recompile()
+
+
+def _arguments(node: Node) -> Dict[str, Any]:
+    arguments = node.target._schema.arguments
+    values = {a.name: a.default_value for a in arguments}
+    values.update(zip((a.name for a in arguments), node.args))
+    values.update(node.kwargs)
+    return values
+
+
+def _sdpa_computes(node: Node, values: Dict[str, Any]) -> bool:
+    if not all(
+        user.target is operator.getitem and (user.args[1] == 0 or not user.users)
+        for user in node.users
+    ):
+        return False
+    if not values["is_causal"]:
+        return True
+    # scaled_dot_product_attention takes no mask together with is_causal.
+    if values.get("attn_bias") is not None:
+        return False
+    # When the query and key lengths differ, flash attention aligns the causal
+    # mask to the bottom right, and scaled_dot_product_attention to the top left.
+    query, key = values["query"].meta["val"], values["key"].meta["val"]
+    return node.target != _FLASH or statically_known_true(
+        query.shape[-2] == key.shape[-2]
+    )
+
+
+def _replace_with_sdpa(
+    module: GraphModule,
+    node: Node,
+    values: Dict[str, Any],
+    signature: Optional[ExportGraphSignature],
+) -> None:
+    query, key = values["query"], values["key"]
+    with module.graph.inserting_before(node):
+        sdpa = module.graph.call_function(
+            exir_ops.edge.aten.scaled_dot_product_attention.default,
+            args=(
+                query,
+                key,
+                values["value"],
+                # Only the flash variant has no bias.
+                values.get("attn_bias"),
+                values["dropout_p"],
+                values["is_causal"],
+            ),
+            kwargs={
+                "scale": values["scale"],
+                # The fused variants that accept fewer key heads than query
+                # heads group them as enable_gqa does.
+                "enable_gqa": not statically_known_true(
+                    query.meta["val"].shape[-3] == key.meta["val"].shape[-3]
+                ),
+            },
+        )
+    sdpa.meta = {**node.meta, "val": node.meta["val"][0]}
+    for user in list(node.users):
+        user.replace_all_uses_with(sdpa)
+        if signature is not None:
+            signature.replace_all_uses(user.name, sdpa.name)
+        module.graph.erase_node(user)
+    module.graph.erase_node(node)
diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py
index 7298b6bc037..bc1e66cf8fe 100644
--- a/backends/cuda/tests/test_cuda_export.py
+++ b/backends/cuda/tests/test_cuda_export.py
@@ -844,6 +844,58 @@ def test_triton_kernel_mode_off(self):
             "SDPA kernel export with triton_kernel_mode=OFF failed",
         )
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_fused_sdpa(self):
+        """Decomposing attention on CUDA, as Torch-TensorRT does, gives fused
+        variants that the backend has no kernel for."""
+        from executorch.backends.cuda.passes.replace_fused_sdpa import (
+            ReplaceFusedSdpaWithSdpaPass,
+        )
+        from executorch.runtime import Runtime
+
+        class FusedSdpa(torch.nn.Module):
+            def forward(self, q, k, v, bias):
+                # Fewer key heads than query heads.
+                flash = torch.ops.aten._scaled_dot_product_flash_attention(
+                    q, k[:, :2], v[:, :2], is_causal=True, scale=0.2
+                )[0]
+                efficient = torch.ops.aten._scaled_dot_product_efficient_attention(
+                    q, k, v, bias, False
+                )[0]
+                return flash, efficient
+
+        model = FusedSdpa()
+        inputs = [torch.randn(1, 4, 16, 64, dtype=torch.bfloat16) for _ in range(3)]
+        inputs.append(torch.randn(1, 4, 16, 16, dtype=torch.bfloat16))
+        edge = self._export_to_cuda_with_lower(model, tuple(inputs))
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.pte")
+            with open(path, "wb") as f:
+                edge.to_executorch().write_to_file(f)
+            method = Runtime.get().load_program(path).load_method("forward")
+            outputs = method.execute(inputs)
+        expected = model(*(t.cuda() for t in inputs))
+        for output, reference in zip(outputs, expected, strict=True):
+            torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2)
+
+        class CausalFlash(torch.nn.Module):
+            def forward(self, q, k, v):
+                return torch.ops.aten._scaled_dot_product_flash_attention(
+                    q, k, v, is_causal=True
+                )[0]
+
+        # With fewer queries than keys, flash aligns the causal mask to the
+        # bottom right, so the pass leaves it.
+        program = _program_for_cuda_passes(
+            CausalFlash(), (inputs[0][:, :, :8], inputs[1], inputs[2])
+        )
+        code = program.graph_module.code
+        ReplaceFusedSdpaWithSdpaPass()(program)
+        self.assertEqual(program.graph_module.code, code)
+
     def test_missing_scatter_kernel_fails_at_export(self):
         """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks."""
 

From 424a7b733ddbbe6d1217acac368812845b0e44c8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 12:27:23 -0800
Subject: [PATCH 4/5] Keep Inductor's decompositions in the antialiased resize
 compile test

test_align_corners_compiles_with_dynamic_sizes compiled the resize with
only the CUDA backend's decomposition table. That replaces Inductor's
own table, so nothing decomposes the arange that the resize uses, and
Inductor falls back to the ATen kernel. When the CI environment variable
is set, Inductor turns that fallback into an error, so the test failed
in CI and passed elsewhere:

    AssertionError: make_fallback(aten.arange.default): a decomposition
    exists, we should switch to it.

The backend itself does not have this problem. AOTI applies the
backend's table to the program and then compiles it with Inductor's own
decompositions. The test now passes both tables, with the backend's
entries on top, the way AOTI sees them. It still fails when the bicubic
decomposition is wrong.
---
 backends/cuda/tests/test_antialiased_resize.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py
index 963444f80e7..09bf9f90b19 100644
--- a/backends/cuda/tests/test_antialiased_resize.py
+++ b/backends/cuda/tests/test_antialiased_resize.py
@@ -155,13 +155,18 @@ def test_keeps_memory_format(self):
 
     def test_align_corners_compiles_with_dynamic_sizes(self):
         from torch._inductor.compile_fx import compile_fx
+        from torch._inductor.decomposition import select_decomp_table
 
         model = Resize("bicubic", align_corners=True)
+        # AOTI lowers with Inductor's own decompositions after the backend's,
+        # and Inductor needs them, for example to lower arange.
+        decompositions = {
+            **select_decomp_table(),
+            **CudaBackend.get_decomposition_table(),
+        }
         compiled = torch.compile(
             model,
-            backend=functools.partial(
-                compile_fx, decompositions=CudaBackend.get_decomposition_table()
-            ),
+            backend=functools.partial(compile_fx, decompositions=decompositions),
             dynamic=True,
             fullgraph=True,
         )

From bbe873c6ab15bbc6f8c9245d3b9231c1a467f709 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 14:16:12 -0800
Subject: [PATCH 5/5] Fix fused attention with one query in the CUDA backend

ReplaceFusedSdpaWithSdpaPass passed scale=None to the
scaled_dot_product_attention it creates. For one query over 256 or more
keys, the Triton pass sends that call to its decode kernel, which takes
a float scale, so the export failed:

    RuntimeError: triton::sdpa_decode_splitk() Expected a value of type
    'float' for argument 'scale' but instead found type 'NoneType'.

The pass now leaves out the arguments that have their default value.

The decode kernel also ignores is_causal, while the efficient and cuDNN
variants align a causal mask to the top left, where one query sees only
the first key. So the pass now keeps every causal call whose query and
key lengths may differ, not only the flash ones. A node that the pass
keeps fails the export later with a missing kernel, so the pass now logs
a warning that says why it kept the node.

test_fused_sdpa now also runs one query over 256 keys, which fails
without this change. test_fused_sdpa_rewrite checks on the CPU which
calls the pass rewrites (cuDNN, causal with equal lengths, inside a
torch.cond branch) and which it keeps (causal with different lengths,
causal with a bias, a used log-sum-exp).
---
 backends/cuda/passes/replace_fused_sdpa.py | 51 ++++++-----
 backends/cuda/tests/test_cuda_export.py    | 98 +++++++++++++++++-----
 2 files changed, 104 insertions(+), 45 deletions(-)

diff --git a/backends/cuda/passes/replace_fused_sdpa.py b/backends/cuda/passes/replace_fused_sdpa.py
index 04b425eb4ec..68743a48ae7 100644
--- a/backends/cuda/passes/replace_fused_sdpa.py
+++ b/backends/cuda/passes/replace_fused_sdpa.py
@@ -26,9 +26,8 @@
 
 logger = logging.getLogger(__name__)
 
-_FLASH = exir_ops.edge.aten._scaled_dot_product_flash_attention.default
 _FUSED_SDPA_OPS = (
-    _FLASH,
+    exir_ops.edge.aten._scaled_dot_product_flash_attention.default,
     exir_ops.edge.aten._scaled_dot_product_efficient_attention.default,
     exir_ops.edge.aten._scaled_dot_product_cudnn_attention.default,
 )
@@ -58,11 +57,16 @@ def __call__(self, exported_program: ExportedProgram) -> None:
                 if node.op != "call_function" or node.target not in _FUSED_SDPA_OPS:
                     continue
                 values = _arguments(node)
-                if _sdpa_computes(node, values):
+                reason = _reason_to_keep(node, values)
+                if reason is None:
                     _replace_with_sdpa(module, node, values, signature)
                     modified = True
                 else:
-                    logger.info(f"Keeping {node.name}: SDPA computes it differently")
+                    logger.warning(
+                        "The CUDA backend cannot lower %s, since %s",
+                        node.format_node(),
+                        reason,
+                    )
             if modified:
                 module.recompile()
 
@@ -75,23 +79,23 @@ def _arguments(node: Node) -> Dict[str, Any]:
     return values
 
 
-def _sdpa_computes(node: Node, values: Dict[str, Any]) -> bool:
+def _reason_to_keep(node: Node, values: Dict[str, Any]) -> Optional[str]:
     if not all(
         user.target is operator.getitem and (user.args[1] == 0 or not user.users)
         for user in node.users
     ):
-        return False
+        return "an output other than the attention output is used"
     if not values["is_causal"]:
-        return True
-    # scaled_dot_product_attention takes no mask together with is_causal.
+        return None
     if values.get("attn_bias") is not None:
-        return False
-    # When the query and key lengths differ, flash attention aligns the causal
-    # mask to the bottom right, and scaled_dot_product_attention to the top left.
+        return "scaled_dot_product_attention takes no mask together with is_causal"
+    # Flash attention aligns the causal mask to the bottom right, the other
+    # variants and scaled_dot_product_attention to the top left, and the Triton
+    # decode kernel ignores it.
     query, key = values["query"].meta["val"], values["key"].meta["val"]
-    return node.target != _FLASH or statically_known_true(
-        query.shape[-2] == key.shape[-2]
-    )
+    if not statically_known_true(query.shape[-2] == key.shape[-2]):
+        return "its causal mask has different query and key lengths"
+    return None
 
 
 def _replace_with_sdpa(
@@ -101,6 +105,16 @@ def _replace_with_sdpa(
     signature: Optional[ExportGraphSignature],
 ) -> None:
     query, key = values["query"], values["key"]
+    kwargs = {}
+    # The Triton kernels take a float scale, so leave out the default.
+    if values["scale"] is not None:
+        kwargs["scale"] = values["scale"]
+    # The fused variants that accept fewer key heads than query heads group
+    # them as enable_gqa does.
+    if not statically_known_true(
+        query.meta["val"].shape[-3] == key.meta["val"].shape[-3]
+    ):
+        kwargs["enable_gqa"] = True
     with module.graph.inserting_before(node):
         sdpa = module.graph.call_function(
             exir_ops.edge.aten.scaled_dot_product_attention.default,
@@ -113,14 +127,7 @@ def _replace_with_sdpa(
                 values["dropout_p"],
                 values["is_causal"],
             ),
-            kwargs={
-                "scale": values["scale"],
-                # The fused variants that accept fewer key heads than query
-                # heads group them as enable_gqa does.
-                "enable_gqa": not statically_known_true(
-                    query.meta["val"].shape[-3] == key.meta["val"].shape[-3]
-                ),
-            },
+            kwargs=kwargs,
         )
     sdpa.meta = {**node.meta, "val": node.meta["val"][0]}
     for user in list(node.users):
diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py
index bc1e66cf8fe..acc823107bc 100644
--- a/backends/cuda/tests/test_cuda_export.py
+++ b/backends/cuda/tests/test_cuda_export.py
@@ -427,6 +427,68 @@ def in_branch(x):
             for name, tensor in {**program.state_dict, **program.constants}.items():
                 self.assertIs(tensor, tensors[name])
 
+    def test_fused_sdpa_rewrite(self):
+        """Fused attention becomes scaled_dot_product_attention only when that
+        computes the same result."""
+        from executorch.backends.cuda.passes.replace_fused_sdpa import (
+            ReplaceFusedSdpaWithSdpaPass,
+        )
+
+        aten = torch.ops.aten
+
+        class Fn(torch.nn.Module):
+            def __init__(self, fn):
+                super().__init__()
+                self.fn = fn
+
+            def forward(self, *args):
+                return self.fn(*args)
+
+        def cudnn(q, k):
+            return aten._scaled_dot_product_cudnn_attention(q, k, k, None, False)[0]
+
+        def causal(q, k, bias=None):
+            return aten._scaled_dot_product_efficient_attention(
+                q, k, k, bias, False, is_causal=True
+            )[0]
+
+        def causal_flash(q, k):
+            return aten._scaled_dot_product_flash_attention(q, k, k, is_causal=True)[0]
+
+        def with_log_sum_exp(q, k):
+            return aten._scaled_dot_product_efficient_attention(q, k, k, None, True)[:2]
+
+        def flash(q, k):
+            return aten._scaled_dot_product_flash_attention(q, k, k)[0]
+
+        def in_branch(q, k):
+            return torch.cond(q.sum() > 0, flash, lambda q, k: q * 2, (q, k))
+
+        q = torch.randn(1, 2, 16, 64)
+        decode = q[:, :, :1].clone()
+        for fn, inputs, rewritten in (
+            (cudnn, (q, q), True),
+            (causal, (q, q), True),
+            (in_branch, (q, q), True),
+            # Flash aligns the causal mask to the bottom right, and the Triton
+            # decode kernel ignores it.
+            (causal, (decode, q), False),
+            (causal_flash, (decode, q), False),
+            (causal, (q, q, torch.randn(1, 2, 16, 16)), False),
+            (with_log_sum_exp, (q, q), False),
+        ):
+            with self.subTest(fn=fn.__name__, query=inputs[0].shape):
+                program = _program_for_cuda_passes(Fn(fn), inputs)
+                ReplaceFusedSdpaWithSdpaPass()(program)
+                fused = [
+                    node
+                    for module in program.graph_module.modules()
+                    if isinstance(module, torch.fx.GraphModule)
+                    for node in module.graph.nodes
+                    if "aten._scaled_dot_product_" in str(node.target)
+                ]
+                self.assertEqual(not fused, rewritten)
+
 
 class TestCudaExport(unittest.TestCase):
     """Test CUDA export functionality for various operations using to_edge_transform_and_lower."""
@@ -851,25 +913,30 @@ def test_triton_kernel_mode_off(self):
     def test_fused_sdpa(self):
         """Decomposing attention on CUDA, as Torch-TensorRT does, gives fused
         variants that the backend has no kernel for."""
-        from executorch.backends.cuda.passes.replace_fused_sdpa import (
-            ReplaceFusedSdpaWithSdpaPass,
-        )
         from executorch.runtime import Runtime
 
         class FusedSdpa(torch.nn.Module):
             def forward(self, q, k, v, bias):
                 # Fewer key heads than query heads.
                 flash = torch.ops.aten._scaled_dot_product_flash_attention(
-                    q, k[:, :2], v[:, :2], is_causal=True, scale=0.2
+                    q, k[:, :2, :16], v[:, :2, :16], is_causal=True, scale=0.2
                 )[0]
                 efficient = torch.ops.aten._scaled_dot_product_efficient_attention(
-                    q, k, v, bias, False
+                    q, k[:, :, :16], v[:, :, :16], bias, False
+                )[0]
+                # One query over many keys runs in the Triton decode kernel.
+                decode = torch.ops.aten._scaled_dot_product_flash_attention(
+                    q[:, :, :1], k, v
                 )[0]
-                return flash, efficient
+                return flash, efficient, decode
 
         model = FusedSdpa()
-        inputs = [torch.randn(1, 4, 16, 64, dtype=torch.bfloat16) for _ in range(3)]
-        inputs.append(torch.randn(1, 4, 16, 16, dtype=torch.bfloat16))
+        inputs = [
+            torch.randn(1, 4, 16, 64, dtype=torch.bfloat16),
+            torch.randn(1, 4, 256, 64, dtype=torch.bfloat16),
+            torch.randn(1, 4, 256, 64, dtype=torch.bfloat16),
+            torch.randn(1, 4, 16, 16, dtype=torch.bfloat16),
+        ]
         edge = self._export_to_cuda_with_lower(model, tuple(inputs))
         with tempfile.TemporaryDirectory() as directory:
             path = os.path.join(directory, "model.pte")
@@ -881,21 +948,6 @@ def forward(self, q, k, v, bias):
         for output, reference in zip(outputs, expected, strict=True):
             torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2)
 
-        class CausalFlash(torch.nn.Module):
-            def forward(self, q, k, v):
-                return torch.ops.aten._scaled_dot_product_flash_attention(
-                    q, k, v, is_causal=True
-                )[0]
-
-        # With fewer queries than keys, flash aligns the causal mask to the
-        # bottom right, so the pass leaves it.
-        program = _program_for_cuda_passes(
-            CausalFlash(), (inputs[0][:, :, :8], inputs[1], inputs[2])
-        )
-        code = program.graph_module.code
-        ReplaceFusedSdpaWithSdpaPass()(program)
-        self.assertEqual(program.graph_module.code, code)
-
     def test_missing_scatter_kernel_fails_at_export(self):
         """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks."""