From aca1ebd872cec956058b061a676b568041594c89 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 20:33:47 -0700 Subject: [PATCH 1/7] 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, complex values inside a torch.cond branch or complex32, it leaves the program as it is and logs why. Complex128 math becomes float64 math, which the CUDA runtime supports. The rotary embeddings of LeRobot's LingBot-VA and Fast-WAM policies are complex128. Two view_as_real outputs of one complex tensor become one output in two slots, and the CUDA runtime deletes such a handle only once. test_complex_rotary_embedding lowers a rotary embedding with a complex constant, polar with a magnitude other than one, views, a split, a concatenation and view_as_real outputs, one of them returned twice, runs it with the runtime and compares it with eager, in float32 and bfloat16, and with complex128 math. The tests in passes/tests/test_replace_complex_with_real.py work on the CPU, with the program prepared the way the CUDA backend's passes get it. One checks a rotary embedding with a complex constant and parameter, a gather by position, views, repeats, splits, joins and squeezes of many kinds and a dynamic length against eager, after the retrace that the backend runs when its decomposition table applies. One checks that the programs listed above, a channels-last copy, a cast to real and a mutated complex buffer stay as they are. One checks that the pass stays registered. They fail without this change. --- backends/aoti/aoti_backend.py | 6 +- backends/cuda/BUCK | 1 + backends/cuda/cuda_backend.py | 4 + .../cuda/passes/replace_complex_with_real.py | 382 ++++++++++++++++++ .../tests/test_replace_complex_with_real.py | 121 ++++++ backends/cuda/tests/test_cuda_export.py | 73 ++++ 6 files changed, 584 insertions(+), 3 deletions(-) create mode 100644 backends/cuda/passes/replace_complex_with_real.py create mode 100644 backends/cuda/passes/tests/test_replace_complex_with_real.py diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index a069db4c1fc..785d89d9c9f 100644 --- a/backends/aoti/aoti_backend.py +++ b/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 ): diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 7e8d4c88675..0b924242754 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/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", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 3f368cdd37e..b35a6373dd4 100644 --- a/backends/cuda/cuda_backend.py +++ b/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(), 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..c9d71812b45 --- /dev/null +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -0,0 +1,382 @@ +# 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. It does not rewrite control flow branches, such +as those of torch.cond, or complex32 values. Otherwise it leaves the program as +it is and logs why, 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 _reason_to_keep( + exported_program: ExportedProgram, complex_nodes: Set[Node], affected: Set[Node] +) -> Optional[str]: + """Why the pass must leave the program as it is, if it must.""" + 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 f"it does not rewrite control flow branches: {node.format_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) + if not supported: + return f"it cannot rewrite {node.format_node()}" + # PyTorch computes complex32 math in float32, and its pair would + # compute in float16. + if any(val.dtype == torch.complex32 for val in _complex_vals(node)): + return f"it does not rewrite complex32: {node.format_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 + } + reason = _reason_to_keep(exported_program, complex_nodes, affected) + if reason is not None: + logger.warning( + "The CUDA backend cannot lower the complex numbers in this " + "program, since %s", + reason, + ) + 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/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py new file mode 100644 index 00000000000..adf06e70c68 --- /dev/null +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -0,0 +1,121 @@ +# 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. + +import unittest + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.passes.replace_complex_with_real import ( + ReplaceComplexWithRealPass, +) +from executorch.backends.cuda.passes.tests.test_replace_fused_sdpa import ( + _Fn, + _program_for_cuda_passes, +) + + +class TestReplaceComplexWithRealPass(unittest.TestCase): + def test_registered(self): + """The rewrite must stay in the CUDA backend's passes.""" + self.assertTrue( + any( + isinstance(p, ReplaceComplexWithRealPass) + for p in CudaBackend.get_custom_passes([]) + ) + ) + + def test_rewrite(self): + """The rewrite of complex numbers gives eager's results.""" + 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.repeat(2, 1).select(-2, 1).detach() + # The cast sets the precision of everything after it. + rotated = scaled * f.to(torch.complex128) * phase + first, second = rotated.chunk(2, dim=-1) + second = torch.cat(second.split([1, 2], dim=-1), -1) + joined = torch.cat([first, second], -1) + return ( + torch.view_as_real(joined.squeeze(-2).clone()), + torch.view_as_real(joined.unsqueeze(0).squeeze((0, -2))), + ) + + model = Rotate() + scale = torch.randn(3, 2) + x = torch.randn(5, 1, 12) + 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)) + + def test_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.""" + + 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()) + + 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]) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 7d1a0e91fea..5367b31d0c4 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -1000,6 +1000,79 @@ def forward(self, x: torch.Tensor) -> Tuple[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, precision): + super().__init__() + self.precision = precision + angles = torch.outer(torch.arange(4.0), torch.rand(12)).to(precision) + 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(angles.abs(), angles) + pairs = torch.view_as_complex( + x.to(self.precision).unflatten(-1, (-1, 2)) + ) + rotated = torch.view_as_real(pairs * grid * spin.unsqueeze(1) * 0.5) + # A view_as_real can also be an output of the program, and two + # of one complex tensor become one output in two slots. + return ( + rotated.flatten(-2).type_as(x), + torch.view_as_real(spin), + torch.view_as_real(spin), + ) + + angles = torch.randn(4, 12) + # Rotary embeddings are often computed in float64, so in complex128. + for dtype, precision in ( + (torch.float32, torch.float32), + (torch.bfloat16, torch.float32), + (torch.float32, torch.float64), + ): + with self.subTest(dtype=dtype, precision=precision): + model = ComplexRope(precision) + 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 9b3a7a4c0708e88d4a15d33573b95b43ca879e89 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 07:11:24 -0800 Subject: [PATCH 2/7] Rewrite Torch-TensorRT's _reshape_copy in the CUDA complex pass Torch-TensorRT turns every view into aten._reshape_copy, and the ops it does not convert reach the CUDA backend in that form. So a rotary embedding that views its complex frequencies, as Llama does with freqs_cis.view(...), reaches ReplaceComplexWithRealPass as a _reshape_copy of a complex tensor. The pass had no rewrite for it, left the program as it was, and lowering failed: RuntimeError: Method forward missing fallback kernels (3 total): - aoti_torch_cuda_reshape - aten::view_as_complex_copy - aten::view_as_real_copy _reshape_copy acts on the real and imaginary parts alike, as view_copy does, so it now keeps its op, with the size of the pair. test_rewrite now calls _reshape_copy on complex frequencies. --- backends/cuda/passes/replace_complex_with_real.py | 2 ++ backends/cuda/passes/tests/test_replace_complex_with_real.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py index c9d71812b45..62241b59616 100644 --- a/backends/cuda/passes/replace_complex_with_real.py +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -81,6 +81,8 @@ def _size(size: List[Any], ndim: int) -> List[Any]: # 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]], ...]] = { + # Torch-TensorRT turns every view into _reshape_copy. + _aten._reshape_copy.default: ((1, _size),), _aten.alias_copy.default: (), _aten.cat.default: ((1, _dim),), _aten.clone.default: (), diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py index adf06e70c68..c6e14178b00 100644 --- a/backends/cuda/passes/tests/test_replace_complex_with_real.py +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -42,7 +42,9 @@ def __init__(self): 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) + f = f.unsqueeze(-2).expand(-1, 2, -1).flatten(-2) + # Torch-TensorRT turns views into _reshape_copy. + f = torch.ops.aten._reshape_copy(f, [-1, 1, 6]) 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. From 4ab701ec7c9653d661f6bc49403ed18a0fdda46a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 15:09:41 -0800 Subject: [PATCH 3/7] Lay out the values of view copies in the CUDA complex pass as copies ReplaceComplexWithRealPass computes the value of a new view copy, such as permute_copy, with its view, because the fake kernels of some view copies make dynamic sizes static. That value kept the strides of the view, while the copy returns a contiguous tensor. Later passes read these strides: ReplaceMaskedCopyWithWherePass could then accept a rewrite that changes the layout of its result, which an as_strided after it reads wrongly. The pass now lays the value out as the copy's would be. test_rewrite checks that the view copies have contiguous values. --- backends/cuda/passes/replace_complex_with_real.py | 15 ++++++++++++--- .../tests/test_replace_complex_with_real.py | 5 +++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py index 62241b59616..086b9c0d4eb 100644 --- a/backends/cuda/passes/replace_complex_with_real.py +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -35,6 +35,7 @@ from typing import Any, Callable, Dict, List, Optional, Set, Tuple import torch +import torch.utils._pytree as pytree from executorch.exir.dialects._ops import ops as exir_ops from torch.export import ExportedProgram from torch.fx import Graph, Node @@ -121,8 +122,9 @@ def _size(size: List[Any], ndim: int) -> List[Any]: # 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. +# view, laid out as the copy's would be. 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, @@ -171,7 +173,14 @@ def call(self, target: Any, *args: Any, **kwargs: Any) -> Node: 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) + value = _VIEWS.get(target, target)(*fake_args, **fake_kwargs) + if target in _VIEWS: + value = pytree.tree_map_only( + torch.Tensor, + lambda t: t.clone(memory_format=torch.contiguous_format), + value, + ) + node.meta["val"] = value return node def parts(self, value: Node) -> List[Node]: diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py index c6e14178b00..a6ba9486d51 100644 --- a/backends/cuda/passes/tests/test_replace_complex_with_real.py +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -7,8 +7,10 @@ import unittest import torch +import torch.utils._pytree as pytree from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.passes.replace_complex_with_real import ( + _VIEWS, ReplaceComplexWithRealPass, ) from executorch.backends.cuda.passes.tests.test_replace_fused_sdpa import ( @@ -74,6 +76,9 @@ def forward(self, x, positions, scale): for node in program.graph.nodes: val = node.meta.get("val") self.assertFalse(isinstance(val, torch.Tensor) and val.is_complex()) + if node.target in _VIEWS: + # A view copy returns new tensors laid out in order. + self.assertTrue(all(t.is_contiguous() for t in pytree.tree_leaves(val))) # 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])): From 536addc9455a60e149f393e56efb978d1ebdf101 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 07:13:57 -0800 Subject: [PATCH 4/7] Store a complex tensor once in the CUDA complex pass When another partitioner takes ops between two that read the same complex buffer, the CUDA partition can get a placeholder for each read, all for one buffer. ReplaceComplexWithRealPass stored the tensor as its real pair once per placeholder, so the second time it called view_as_real on the stored real pair, and export stopped: RuntimeError: view_as_real is only supported for complex tensors The pass now stores a tensor only while it is still complex, and gives every placeholder of it the real value. test_tensor_with_two_placeholders gives a buffer two placeholders, as partitioning does, and checks the rewritten program against eager. --- .../cuda/passes/replace_complex_with_real.py | 11 ++++--- .../tests/test_replace_complex_with_real.py | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py index 086b9c0d4eb..8dbb86a6d94 100644 --- a/backends/cuda/passes/replace_complex_with_real.py +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -309,10 +309,13 @@ def _store_as_real(exported_program: ExportedProgram, node: Node) -> None: 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 + # Partitioning can give one tensor several placeholders, and the first of + # them stores it. + if tensor.is_complex(): + 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()) diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py index a6ba9486d51..8797209ae30 100644 --- a/backends/cuda/passes/tests/test_replace_complex_with_real.py +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -17,6 +17,9 @@ _Fn, _program_for_cuda_passes, ) +from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import ( + duplicate_constant_node, +) class TestReplaceComplexWithRealPass(unittest.TestCase): @@ -85,6 +88,32 @@ def forward(self, x, positions, scale): inputs = (x[: len(positions)], positions, scale) torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + def test_tensor_with_two_placeholders(self): + """Partitioning gives a tensor a placeholder per op that reads it when + another backend takes ops between those, and the pass stores the + tensor once.""" + + class Rotate(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("freqs", torch.polar(torch.ones(3), torch.rand(3))) + + def forward(self, x, y): + return tuple( + torch.view_as_real(torch.view_as_complex(t) * self.freqs) + for t in (x, y) + ) + + model = Rotate() + inputs = (torch.randn(3, 2), torch.randn(3, 2)) + program = _program_for_cuda_passes(model, inputs) + duplicate_constant_node(program, "b_freqs") + ReplaceComplexWithRealPass()(program) + for node in program.graph.nodes: + val = node.meta.get("val") + self.assertFalse(isinstance(val, torch.Tensor) and val.is_complex()) + torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + def test_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 2bc1dbf10ab1a7b166c876f15996ef9e3007a569 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 07:27:43 -0800 Subject: [PATCH 5/7] Make the docstrings of the CUDA complex pass match its code The module docstring said "Otherwise" right after the sentence about control flow and complex32, so it read as if the pass left every other program alone. It also named only constants as stored in the real layout, while parameters and buffers are too, and left out the casts from real to complex. The docstring of _constant_name promised a name only for tensors the graph never writes, while only buffers are checked. --- .../cuda/passes/replace_complex_with_real.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py index 8dbb86a6d94..baeb1183474 100644 --- a/backends/cuda/passes/replace_complex_with_real.py +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -20,14 +20,16 @@ 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 complex tensor. A cast of a real tensor to complex gets zero imaginary +parts. A complex constant, parameter or buffer is stored in the real layout +too. 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. It does not rewrite control flow branches, such -as those of torch.cond, or complex32 values. Otherwise it leaves the program as -it is and logs why, and the backend reports the ops it has no kernels for. +writes a complex tensor in it. Otherwise it leaves the program as it is and +logs why, and the backend reports the ops it has no kernels for. It never +rewrites control flow branches, such as those of torch.cond, or complex32 +values. """ import logging @@ -286,8 +288,8 @@ def _can_rewrite(node: Node, complex_nodes: Set[Node]) -> bool: 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.""" + """The name of the constant, parameter or buffer that ``node`` reads, + unless it is a buffer that the graph writes.""" signature = exported_program.graph_signature for inputs in ( signature.inputs_to_lifted_tensor_constants, From 2eb93e7c0107124fe5a18e4653f455c333e8826f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 07:37:43 -0800 Subject: [PATCH 6/7] Refuse copies and clones inside a CUDA graph capture A CUDA graph capture records kernels without running them, and it cannot record a synchronous copy. The .item() and mask index shims check for an active capture and return NotSupported, so the backend runs the call again without a graph. aoti_torch_copy_, aoti_torch_clone and aoti_torch_clone_preserve_strides had no such check. They went straight to a synchronous copy, which fails inside a capture, and the failed copy aborts the process. AOTI returns a part of a constant through aoti_torch_clone, so a method as small as def forward(self, x): return self.table[1], x * 2 aborted at its first capture when CUDA graphs were on: CUDA error: operation would make the legacy stream depend on a capturing blocking stream These shims now make the same check first, through one helper that the .item() shims share, and their docs say when they return NotSupported. An empty tensor copies nothing, so it still captures as before. test_aoti_torch_copy_ checks that copies and clones of CUDA tensors refuse during a capture, and that a clone in host memory and an empty copy do not. test_constant_copy_runs_without_a_graph runs the method above with CUDA graphs on. --- .../tests/test_replace_complex_with_real.py | 4 +- backends/cuda/runtime/shims/memory.cpp | 50 +++++++----- backends/cuda/runtime/shims/memory.h | 11 ++- .../shims/tests/test_aoti_torch_copy_.cpp | 78 +++++++++++++++++++ .../tests/test_output_stride_rearrange.py | 25 ++++++ 5 files changed, 146 insertions(+), 22 deletions(-) diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py index 8797209ae30..793fe16555c 100644 --- a/backends/cuda/passes/tests/test_replace_complex_with_real.py +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -107,7 +107,9 @@ def forward(self, x, y): model = Rotate() inputs = (torch.randn(3, 2), torch.randn(3, 2)) program = _program_for_cuda_passes(model, inputs) - duplicate_constant_node(program, "b_freqs") + self.assertEqual( + duplicate_constant_node(program, "b_freqs"), {"b_freqs_copy_0"} + ) ReplaceComplexWithRealPass()(program) for node in program.graph.nodes: val = node.meta.get("val") diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index a9a39a7c7b1..f36b283b4c0 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -46,6 +46,29 @@ using c10::ScalarType; namespace { +// A graph capture cannot record a synchronous copy to or from device memory, +// and such a copy aborts the process when it fails, so a shim that makes one +// refuses first. Info, not Error: the backend then retries the call without a +// graph. An empty tensor copies nothing, so it needs no check. +AOTITorchError refuse_inside_capture( + const SlimTensor* tensor, + const char* what) { +#if defined(CUDA_AVAILABLE) + if (tensor->is_cuda() && tensor->numel() > 0) { + auto stream = getCurrentCUDAStream(tensor->device_index()); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaStreamIsCapturing(stream.get(), &capture)); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG(Info, "%s: not supported inside a CUDA graph capture", what); + return Error::NotSupported; + } + } +#endif + return Error::Ok; +} + // Reads the one element as T. Dispatch on the dtype the tensor actually holds: // item() copies sizeof(T) bytes and does not check, so asking for the wrong // width reads past a one-element allocation. @@ -67,25 +90,8 @@ AOTITorchError narrow_to(From value, To* ret_value) { template AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { -#if defined(CUDA_AVAILABLE) - // A graph capture cannot record the copy to the host below, and the copy - // aborts the process when it fails, so refuse first. Info, not Error: the - // backend then retries the call without a graph. - if (tensor->is_cuda()) { - auto stream = getCurrentCUDAStream(tensor->device_index()); - ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); - cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; - ET_CUDA_CHECK_OR_RETURN_ERROR( - cudaStreamIsCapturing(stream.get(), &capture)); - if (capture != cudaStreamCaptureStatusNone) { - ET_LOG( - Info, - "reading a single element: a device value cannot be read inside a " - "CUDA graph capture"); - return Error::NotSupported; - } - } -#endif + ET_CHECK_OK_OR_RETURN_ERROR( + refuse_inside_capture(tensor, "reading a single element")); switch (tensor->dtype()) { case ScalarType::Byte: return narrow_to(tensor->item(), ret_value); @@ -352,6 +358,8 @@ AOTITorchError aoti_torch_clone_preserve_strides( ret_new_tensor != nullptr, InvalidArgument, "aoti_torch_clone_preserve_strides: ret_new_tensor is null"); + ET_CHECK_OK_OR_RETURN_ERROR( + refuse_inside_capture(self, "aoti_torch_clone_preserve_strides")); *ret_new_tensor = new SlimTensor(self->clone()); return Error::Ok; @@ -365,6 +373,7 @@ AOTITorchError aoti_torch_clone(SlimTensor* self, SlimTensor** ret_new_tensor) { ret_new_tensor != nullptr, InvalidArgument, "aoti_torch_clone: ret_new_tensor is null"); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(self, "aoti_torch_clone")); *ret_new_tensor = new SlimTensor(self->clone()); return Error::Ok; @@ -437,6 +446,9 @@ aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) { ET_CHECK_OR_RETURN_ERROR( src != nullptr, InvalidArgument, "aoti_torch_copy_: src is null"); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(self, "aoti_torch_copy_")); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(src, "aoti_torch_copy_")); + // SlimTensor::copy_() handles: // - Same numel validation // - Same dtype validation diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index ef51ff5eef0..4c4d67b7a53 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -150,6 +150,10 @@ aoti_torch_new_tensor_handle(SlimTensor* orig_handle, SlimTensor** new_handle); /** * Clone APIs. + * + * A clone of a CUDA tensor is a synchronous copy, which a CUDA graph capture + * cannot record, so both return Error::NotSupported while the current stream + * is capturing. */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_clone_preserve_strides( SlimTensor* self, @@ -207,7 +211,9 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_as_strided( * @param self Destination tensor (must not be null) * @param src Source tensor to copy from (must not be null) * @param non_blocking If true, the copy may be asynchronous (currently ignored) - * @return AOTITorchError error code (Error::Ok on success) + * @return Error::Ok on success, Error::NotSupported if either tensor is a CUDA + * tensor and the current stream is capturing a CUDA graph, which cannot record + * the synchronous copy */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking); @@ -231,7 +237,8 @@ aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); * * @param tensor Single-element tensor (must not be null) * @param ret_value Output parameter for the extracted value - * @return AOTITorchError error code (Error::Ok on success) + * @return Error::Ok on success, Error::NotSupported for a CUDA tensor while the + * current stream is capturing a CUDA graph, which cannot record the copy */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_uint8(SlimTensor* tensor, uint8_t* ret_value); diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp index 392a3e29244..df67f002184 100644 --- a/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -449,6 +450,83 @@ TEST_F(AOTITorchCopySlimTest, CudaToCpu) { EXPECT_EQ(aoti_torch_delete_tensor_object(dst), Error::Ok); } +TEST_F(AOTITorchCopySlimTest, RefusedInsideGraphCapture) { + if (!isCudaAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + + std::vector sizes = {2, 3}; + Tensor* device = createTestTensor( + sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* other_device = createTestTensor( + sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* host = createTestTensor(sizes); + std::vector empty_sizes = {0, 3}; + Tensor* empty_device = createTestTensor( + empty_sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* other_empty_device = createTestTensor( + empty_sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + ASSERT_NE(device, nullptr); + ASSERT_NE(other_device, nullptr); + ASSERT_NE(host, nullptr); + ASSERT_NE(empty_device, nullptr); + ASSERT_NE(other_empty_device, nullptr); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ(setCurrentCUDAStream(stream, 0), Error::Ok); + ASSERT_EQ( + cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), + cudaSuccess); + // A capture cannot record these synchronous copies, so each one refuses + // instead of failing inside the copy. + EXPECT_EQ(aoti_torch_copy_(host, device, 0), Error::NotSupported); + EXPECT_EQ(aoti_torch_copy_(device, host, 0), Error::NotSupported); + EXPECT_EQ(aoti_torch_copy_(other_device, device, 0), Error::NotSupported); + Tensor* clone = nullptr; + EXPECT_EQ(aoti_torch_clone(device, &clone), Error::NotSupported); + EXPECT_EQ( + aoti_torch_clone_preserve_strides(device, &clone), Error::NotSupported); + // A copy in host memory does not touch the stream, and an empty one copies + // nothing. + EXPECT_EQ(aoti_torch_clone(host, &clone), Error::Ok); + if (clone != nullptr) { + EXPECT_EQ(aoti_torch_delete_tensor_object(clone), Error::Ok); + } + EXPECT_EQ(aoti_torch_copy_(empty_device, other_empty_device, 0), Error::Ok); + cudaGraph_t graph = nullptr; + // Nothing reached the stream, so the capture is still valid. + EXPECT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + if (graph != nullptr) { + EXPECT_EQ(cudaGraphDestroy(graph), cudaSuccess); + } + EXPECT_EQ(aoti_torch_copy_(host, device, 0), Error::Ok); + + clearCurrentCUDAStream(0); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); + EXPECT_EQ(aoti_torch_delete_tensor_object(device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(other_device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(host), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(empty_device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(other_empty_device), Error::Ok); +} + // ============================================================================ // Non-blocking Tests // ============================================================================ diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index a53156cd3f0..176acd089e5 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -316,6 +316,17 @@ def forward(self, x, mask): return self.calls.clone(), x[mask] +class _ConstantRow(nn.Module): + """AOTI returns a part of a constant through aoti_torch_clone.""" + + def __init__(self): + super().__init__() + self.register_buffer("table", torch.randn(3, 4)) + + def forward(self, x): + return self.table[1], x * 2 + + class _HostFlag(nn.Module): """A torch.cond on a flag the program keeps in host memory and updates on every call.""" @@ -531,6 +542,20 @@ def test_mask_retry_updates_state_once(self): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, torch.tensor([i + 1.0])) + def test_constant_copy_runs_without_a_graph(self): + # A capture cannot record the copy that returns a part of a constant, + # so the capture call has to run without a graph. + torch.manual_seed(0) + module = _ConstantRow().eval() + inputs = [torch.randn(4) for _ in range(6)] + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, (inputs[0],), inputs) + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) + for i, out in enumerate(outputs): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module.table[1]) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. From 199a90e267f2687ac9b57ce461c2606cf027b000 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 07:37:43 -0800 Subject: [PATCH 7/7] Compare the varlen_attn decomposition on ROCm where windows agree test_varlen_attn compares the varlen_attn decomposition with the op. On ROCm the op runs AOTriton, which lines an explicit window up with the start of each sequence, while the decomposition follows Flash attention on CUDA, which lines it up with the end. So the test compared nothing on ROCm. Both place a window the same way when the query and key sequences line up, so the test now skips on ROCm only the cases with an explicit window over query and key sequences that differ, and compares the rest there too. This was not run on ROCm. The only ROCm run of the comparison failed at its first case, a window over query sequences shorter than their key sequences, and never reached the others. On CUDA the test compares every case, as before. --- backends/cuda/tests/test_cuda_export.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 5367b31d0c4..1e9489479bc 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -556,13 +556,21 @@ def forward(self, q, k, v, cu_seq): (big, big, v, cu_seq, cu_seq, 16, 16), ] table = CudaBackend.get_decomposition_table() - # On ROCm, varlen_attn runs AOTriton, which lines an explicit window up - # with the start of each sequence. The decomposition follows Flash - # attention on CUDA, which lines it up with the end, so compare on CUDA. - if torch.version.hip is None: - for args in cases: - for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): - torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + for args in cases: + # On ROCm, varlen_attn runs AOTriton, which lines an explicit window + # up with the start of each sequence. The decomposition follows + # Flash attention on CUDA, which lines it up with the end. The two + # agree when the query and key sequences line up, as in + # self-attention. + window = args[9] if len(args) > 9 else None + if ( + torch.version.hip is not None + and window is not None + and not torch.equal(args[3], args[4]) + ): + continue + for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) # Flash attention versions differ on this window, since a query sequence # can be longer than max_k + 1. with self.assertRaisesRegex(NotImplementedError, "window"):