[shoumikhin]

What is wrong today

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 cannot 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, so it calls ATen kernels, and this backend ships none of them. Torch-TensorRT leaves these ops to the CUDA backend, so its programs stop at the same error.

What this change does

A new pass, ReplaceComplexWithRealPass, stores 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:

• polar becomes cosine and sine times the magnitude.
• A complex product becomes four real products.
• view_as_complex and view_as_real go away.
• Views, slices, splits, indexing, repeats, transposes, joins and casts keep their op, with the dims and dtype of the pair.
• Complex constants, buffers and parameters are stored as their real view.

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

The pass rewrites a program only when it can rewrite every op on a complex tensor. 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.

What was tested

With torch 2.14 on an A100:

• test_complex_rotary_embedding runs a lowered rotary embedding with a complex constant, polar with a magnitude other than one, and a view_as_real output returned twice, and compares with eager in float32 and bfloat16, and with complex128 math.
• test_replace_complex_with_real.py runs the rewritten program on the CPU against eager, with complex128 math, a dynamic length, repeats and squeezes. It also checks that the programs listed above stay as they are, and that the pass stays registered.

They fail without this change. The other CUDA tests in the same CI job pass.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24403

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 10 Pending, 1 Unrelated Failure

As of commit d90c110 with merge base e2bcb9d ():

NEW FAILURES - The following jobs have failed:

• pull / android / run-emulator (gh)
  Warning: An error occurred while preparing SDK package Android Emulator: Error on ZipFile unknown archive.
• pull / test-arm-backend-vkml (test_pytest_ops_vkml) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 56

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 27m 56s —— View job

---

This PR lowers complex-number rotary embeddings in the CUDA backend and stops repeated AOTI output handles from being deleted twice. Nothing blocks merging, but there are six minor issues: a handle leak in the CUDA-graph capture step, rewritten constants missing the autotune out-of-memory offload, helper code that duplicates backends/transforms/utils.py, and three test issues (no GPU test of dynamic shapes, a weak check in the unsupported-program test, and test placement).

Detailed review

PR Review: #24403

Summary

The PR adds ReplaceComplexWithRealPass, which stores each complex tensor as a real tensor with a trailing dim of 2 before Inductor runs. It also makes the CUDA runtime give each repeated output slot its own handle. All findings below are minor and none blocks merging. Three reviewers covered all 7 changed files.

Code Quality

• The CUDA-graph capture step leaks the new handles. backends/cuda/runtime/cuda_backend.cpp:1034 creates an ExecuTorch-owned SlimTensor for each repeated output slot. The capture branch releases outputs by setting the slots to nullptr (lines 1074, 1089, 1119). That is right for AOTI's own handles, because the graph owns them, but the copies do not belong to AOTI.

  • So each capture of a method that returns one tensor twice leaks one SlimTensor, plus one reference to its storage, per repeated slot. The leak is small and happens once per capture.
  • Deleting the copies at those points only lowers a reference count.
  • No test covers this path. test_same_tensor_returned_twice (test_cuda_export.py:909) runs without CUDA graphs. TestStridedOutputWithCudaGraph._run_with_cuda_graph (backends/cuda/tests/test_output_stride_rearrange.py:304) already runs warmup, capture and replay, but it returns only the first output.

  Fix this →

• _constant_name and _store_as_real duplicate existing helpers. backends/cuda/passes/replace_complex_with_real.py:278-306 looks up parameters, buffers and lifted constants by hand, and writes to the private exported_program._constants / _state_dict.

  • backends/transforms/utils.py:69-128 already has is_param_node, get_param_tensor and set_param_tensor, and the cuda_passes BUCK target already depends on //executorch/backends/transforms:utils.
  • The buffers_to_mutate exclusion (lines 288-290) is also redundant. A mutated complex buffer always passes a complex value to the output node, and _can_rewrite rejects that node at line 256. The test's MutatedBuffer case is therefore declined either way.

  Fix this →

Infrastructure

• Rewritten constants miss the autotune out-of-memory offload. This offload moves the program's GPU tensors to the CPU when an autotune capture runs out of memory.

  • CudaBackend.move_program_to_device records weak references to the program's CUDA tensors (_track_moved_tensors, called at backends/cuda/cuda_backend.py:1283). autotune_launch_params and cuda_graph_timing offload exactly those tensors (lines 1218, 1226).
  • The custom passes run after that step (backends/aoti/aoti_backend.py:296-311). _store_as_real (replace_complex_with_real.py:302-305) then replaces each complex constant with a new .clone().
  • As a result, the tracked tensor drops out and the new real tensor is never tracked. A large complex buffer, such as a long-context frequency table, stays on the GPU when the offload runs.
  • Fix: track the program's tensors again after the custom passes run.

  Fix this →

Testing

• Two behaviors the description reports have no GPU test.

  • Dynamic sequence length. test_complex_rewrite (test_cuda_export.py:319) exports with Dim("length"), but only runs the eager program.module() on the CPU after run_decompositions({}) (line 366). test_complex_rotary_embedding (line 935) uses static shapes. The _VIEWS comment (replace_complex_with_real.py:121-124) says dynamic sizes are where the fake kernels of view copies break, so CI should send dynamic sizes through Inductor/AOTI.
  • view_as_real of one complex tensor returned twice. The pass maps both outputs to one node, and the description gives this as one of the two cases the runtime fix exists for. test_same_tensor_returned_twice only covers y.detach() returned twice.

  Fix this →

• The unsupported-program test cannot detect a partly edited graph. test_complex_rewrite_leaves_unsupported_programs compares program.graph_module.code (test_cuda_export.py:423, 426).

  • FX regenerates that string only in recompile(), and on this path the pass returns before recompile() (replace_complex_with_real.py:384). A regression that edits nodes and then bails out would still pass.
  • The test also does not check placeholder meta["val"], which _store_as_real changes at line 306.
  • Fix: compare str(program.graph) and the placeholder values instead.

  Fix this →

• The pass tests do not follow the per-pass test layout. _program_for_cuda_passes, test_complex_rewrite and test_complex_rewrite_leaves_unsupported_programs (test_cuda_export.py:30, 319, 377) sit in TestCudaBackendCompileOptions.

  • Each sibling pass has its own module in backends/cuda/passes/tests/, for example test_replace_int64_floordiv.py, which has a separate registration test at line 51. CI runs that directory in the same pytest call (.github/workflows/cuda.yml:622).
  • Here the registration check is appended to the end of the numerics test (lines 370-375), so a numerics failure hides it.

  Fix this →

Recommendation

Approve

Nothing here blocks merging. The fixes most worth making before merge are the handle leak in the CUDA-graph capture step and GPU tests for dynamic shapes and for repeated outputs with CUDA graphs.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 24m 19s —— View job

---

Nothing blocks this PR, but ReplaceFusedSdpaWithSdpaPass can silently stop matching eager for causal efficient attention at decode shapes, and the pass has two smaller problems: a hidden log message and test gaps.

Detailed review

PR Review: #24403

Summary

The PR stores complex tensors as pairs of real numbers, turns fused attention variants back into scaled_dot_product_attention, and gives each repeated AOTI output slot its own handle. There are three minor issues and no blockers. Three reviewers covered all 8 changed files.

Code Quality

• Causal efficient and cuDNN attention at decode shapes stop matching eager. _sdpa_computes (backends/cuda/passes/replace_fused_sdpa.py:89-94) compares query and key lengths only for flash. A causal efficient or cuDNN call with L_q != L_kv therefore becomes scaled_dot_product_attention(is_causal=True).

  • ReplaceEdgeOpWithTritonOpPass runs after this pass (backends/cuda/cuda_backend.py:1045-1046). It sends calls with L_q == 1, a static L_kv >= 256 and a power-of-2 head dim to sdpa_decode_splitk (backends/cuda/triton/replacement_pass.py:113-122). That is the only kernel it lets take a non-square causal call (line 181).
  • sdpa_decode_splitk ignores is_causal (backends/cuda/triton/kernels/sdpa.py:1846), so the query attends to every key. Eager efficient attention masks from the top left, so there the query attends only to key 0.
  • Trigger: _scaled_dot_product_efficient_attention(q, k, v, None, False, is_causal=True) with q of shape [1, H, 1, 64] and k and v of shape [1, H, 256, 64].
  • Models that work in eager rarely make this call, because in eager it attends only to key 0. This backend also already lowers a user-written causal SDPA at these shapes the same way (backends/cuda/tests/test_sdpa_splitk_replacement.py:311). Still, the pass says it rewrites only "when that computes the same result", and its comment at lines 89-90 assumes the lowered SDPA masks from the top left.
  • Fix: require statically_known_true(L_q == L_kv) for every causal call, not only flash. Alternatively, pass an explicit top-left mask for non-square causal calls.

  Fix this →

• A kept fused node is logged at a level that is hidden by default, without its reason. replace_fused_sdpa.py:65 logs Keeping ...: SDPA computes it differently with logger.info.

  • A kept node always makes the export fail later with missing fallback kernels, and this line is the only explanation.
  • ReplaceComplexWithRealPass uses logger.warning in the same situation (replace_complex_with_real.py:355).
  • The message does not say which condition kept the node: a used log-sum-exp or debug mask, a causal call with a bias, or a causal flash call with different query and key lengths.

  Fix this →

Testing

• test_fused_sdpa (backends/cuda/tests/test_cuda_export.py:851-897) tests two rewrites and only one of the three conditions that keep a node.

  • Covered: a causal flash call with fewer key heads, a non-causal efficient call with a bias, and a causal flash call with L_q < L_kv, which the pass keeps.
  • Not covered:
    • _scaled_dot_product_cudnn_attention. The pass assumes it masks from the top left and groups fewer key heads the way enable_gqa does (lines 118-119).
    • Any causal efficient call, including the decode shape above.
    • A causal call with a bias, which should be kept.
    • A call whose log-sum-exp is used, which should be kept.
  • The one keep check runs only on CPU, but it sits at the end of the GPU test. So it is skipped on ROCm and never runs when the GPU part fails. A CPU test next to test_complex_rewrite (line 319), built with _program_for_cuda_passes, could cover each case.

  Fix this →

Recommendation

Approve

Nothing blocks merging. The first finding is worth fixing first: it is a one-line change that makes the pass follow its own "same result" rule.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

Claude finished @shoumikhin's task in 30m 12s —— View job

---

Nothing blocks this PR, but if it lands before #24326, complex128 programs that fail at export today will export and then fail to load or abort at run time; the other three findings are minor.

Detailed review

PR Review: #24403

Summary

The PR stores complex tensors as pairs of real numbers, turns fused attention back into scaled_dot_product_attention, and gives each repeated AOTI output slot its own handle. There are no blockers. One finding depends on the merge order with #24326, and three are minor. Three reviewers covered all 9 changed files.

Code Quality

• If this lands before CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326, complex128 programs fail at load or run time instead of at export.

  • The pass computes complex128 as float64 pairs, for example in _copy and _store_as_real (backends/cuda/passes/replace_complex_with_real.py:209, :302). _first_unsupported declines only complex32 (lines 327-332).
  • The runtime on this PR's base has no float64: Double is commented out in backends/aoti/slim/c10/core/ScalarType.h:33.
    • Loading a float64 constant fails. CudaWeightCache::parse returns InvalidProgram for dtype 7 (backends/cuda/runtime/test/test_cuda_weight_cache.cpp:278-284).
    • A float64 buffer that AOTI allocates at run time aborts the process. SlimTensor::check_supportive is an ET_CHECK_MSG (backends/aoti/slim/core/slim_tensor.h:594-599), and backends/aoti/slim/factory/test/test_empty.cpp:252-287 expects that abort.
  • Trigger: a rotary table built in float64, like the Wan complex128 case in the description. Today such a program fails at export with the list of missing kernels.
  • The description says this PR should land after CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups #24326, but nothing enforces that order. If the order cannot be guaranteed, decline complex128 in _first_unsupported until the runtime has float64.

  Fix this →

• Metal still deletes a repeated output handle twice. The PR fixes this only in backends/cuda/runtime/cuda_backend.cpp:1028-1038.

  • MetalBackend is also an AotiBackend (backends/apple/metal/metal_backend.py:21), so it runs the same ReplaceViewCopyWithViewPass (backends/aoti/aoti_backend.py:303) and the same AOTI code generation.
  • Its cleanup guard deletes every gpu_outputs[i] (backends/apple/metal/runtime/metal_backend.cpp:505-511), so return y.detach(), y.detach() most likely double-frees on Metal too.
  • Metal's shims already have aoti_torch_new_tensor_handle (backends/apple/metal/runtime/shims/memory.h), so the same loop ports directly.
  • The bug predates this PR and is outside the diff. I found it by reading the code and did not run it on Metal.

  Fix this →

Testing

• No test covers the repeated output that the complex pass creates.

  • The description gives view_as_real of one complex tensor returned twice as a reason for the runtime fix. The pass turns both outputs into one node (replace_complex_with_real.py:374-382), so AOTI puts one handle in two slots.
  • No test returns one view_as_real twice. test_complex_rotary_embedding returns two different tensors (backends/cuda/tests/test_cuda_export.py:1066).
  • test_same_tensor_returned_twice (line 1013) depends on AOTI giving two separate alias_copy outputs one handle, which is an Inductor detail. If that changes, the test still passes but no longer exercises the fix.
  • Fix: return torch.view_as_real(spin) a second time in test_complex_rotary_embedding. That one test then covers both the pass's signature update for two outputs and the runtime fix.

  Fix this →

• The CPU tests of the new passes don't follow the per-pass test layout.

  • _program_for_cuda_passes, test_complex_rewrite, test_complex_rewrite_leaves_unsupported_programs and test_fused_sdpa_rewrite (test_cuda_export.py:30, 319, 377, 430) sit in TestCudaBackendCompileOptions.
  • Each sibling pass has its own module in backends/cuda/passes/tests/, such as test_replace_int64_floordiv.py, with a separate registration test (line 51). CI runs that directory in the same pytest call (.github/workflows/cuda.yml:622).
  • Here the registration check sits at the end of the numerics test (test_cuda_export.py:370-375), so a numerics failure hides it.

  Fix this →

Recommendation

Approve

Nothing blocks merging as long as this lands after #24326, as the description plans. If that order cannot be guaranteed, decline complex128 first.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

Claude finished @shoumikhin's task in 30m 48s —— View job

---

Nothing blocks this PR, but there are three minor findings: the copy nodes the pass creates carry a view's metadata, a mutated-buffer check never decides anything, and no test compiles a dynamic-length program or covers polar on the CPU.

Detailed review

PR Review: #24403

Summary

The PR adds ReplaceComplexWithRealPass, which stores each complex tensor as a real tensor with a last dim of 2, so that the CUDA backend can lower rotary embeddings written with complex numbers. There are three minor findings and no blockers. Three reviewers covered all 6 changed files.

Code Quality

• New copy nodes get the value of a view, not of a copy. _Rewriter.call (backends/cuda/passes/replace_complex_with_real.py:172) computes the value of each new node whose op is in _VIEWS (lines 124-138), such as permute_copy, transpose_copy or expand_copy, with the matching view op.

  • That value aliases the input and keeps the view's strides: permuted for permute_copy and transpose_copy, zero for expand_copy. The copy op itself returns a new contiguous tensor.
  • Trigger: in torch.view_as_real(z.transpose(1, 2)), the pass replaces the view_as_real_copy node, whose value is contiguous, with a transpose_copy node whose value has transposed strides. Its users see those strides.
  • The passes that run after this one read meta["val"]. ReplaceMaskedCopyWithWherePass._mask_of_copy decides from its strides (backends/cuda/passes/replace_masked_copy_with_where.py:144-146). I have not checked whether this changes any lowering today, so treat the impact as low confidence.
  • No test checks these values. test_rewrite only asserts that none of them is complex (test_replace_complex_with_real.py:72-74), then retraces the program (line 76), which recomputes them.
  • Fix: for the copy targets, make the view's value contiguous, for example with .clone(memory_format=torch.contiguous_format). Check that this keeps the dynamic sizes that the comment at lines 120-123 is about.

  Fix this →

• The mutated-buffer check in _constant_name never decides anything. (replace_complex_with_real.py:287-289)

  • A complex buffer that the graph writes always passes a complex value to the output node. So the output node is in affected, and _can_rewrite rejects it at line 255. The program is declined with or without lines 287-289.
  • As a result, the MutatedBuffer case in test_leaves_unsupported_programs (test_replace_complex_with_real.py:85-92, 113) passes either way. The test never checks which reason the pass logs, for any of its cases.
  • Fix: remove lines 287-289 and "if the graph never writes it" from the docstring (lines 278-279). Then assert the logged reason for each case with assertLogs, so that each case tests the condition it was written for.

  Fix this →

Testing

• No test compiles a rewritten program with a dynamic length, and no CPU test covers polar.

  • test_complex_rotary_embedding (backends/cuda/tests/test_cuda_export.py:970) is the only test of this pass that sends a rewritten program through AOTI and the CUDA runtime. It exports with static shapes, and _export_to_cuda_with_lower (line 357) takes no dynamic_shapes. Rotary embeddings in LLMs are usually exported with a dynamic sequence length.
  • test_rewrite uses a dynamic length, but runs the program only on the CPU, and only after run_decompositions({}) (test_replace_complex_with_real.py:76). The backend runs that retrace only when its decomposition table matches an op in the graph (backends/aoti/aoti_backend.py:337-342). Otherwise it compiles the graph as the pass leaves it.
  • Both torch.polar calls in test_rewrite run in __init__ (lines 38 and 40), so its graph has no polar node. _polar, which builds the rotations, is tested only by the static GPU test. That test is skipped on ROCm (test_cuda_export.py:966-969), although ROCm CI runs the pass tests (.ci/scripts/test-rocm-aoti.sh:157).
  • Fix: add a GPU test that exports a rotary embedding with a dynamic sequence length and runs it at two lengths, for example the Llama form that slices a complex frequency table to the input's length. Call torch.polar in Rotate.forward.

  Fix this →

Recommendation

Approve

Nothing blocks merging, and all three findings are minor. The dynamic-length GPU test is the most useful one to add, because that is how rotary embeddings are usually exported.

[shoumikhin]

Thanks for the review. I checked the first and third on an A100, and the second in the code, before deciding. None of them changes a result today, so I will not push again for them now.

1. Copy nodes with view strides in meta["val"]. You are right: the value of a new permute_copy, transpose_copy or expand_copy has the strides of the view, not of the copy. I lowered and ran torch.view_as_real(z.transpose(1, 2)) through the CUDA backend, and the output matches eager. Inductor traces the graph again, so it does not read these values. The masked-copy pass is the only later pass that makes a decision from strides. A wrong stride can only change whether it rewrites a masked copy as where, and both give the same values. I will make these values contiguous in a follow-up after this lands.

2. The mutated-buffer check in _constant_name. Agreed, it never decides anything, because the output node already makes the pass decline. I will remove it in the same follow-up, and check the logged reason for each case in test_leaves_unsupported_programs.

3. Dynamic length on the GPU. I lowered two Llama-style rotary embeddings with a dynamic sequence length: one calls torch.polar in forward, the other slices a complex buffer. Both ran at lengths 7 and 33 and match eager. In the same follow-up I will add a GPU test like that, and call polar in forward in test_rewrite.