Repository navigation
CUDA backend: lower complex rotary embeddings and free repeated outputs once - #24403
shoumikhin wants to merge 2 commits into
Conversation
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.
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.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24403
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Unclassified FailureAs of commit 9181a61 with merge base 05092c1 ( UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
@claude /pr-review |
|
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 Detailed reviewPR Review: #24403SummaryThe PR adds Code Quality
Infrastructure
Testing
RecommendationApprove 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. |
What is wrong today
Rotary position embeddings are often written with complex numbers (
torch.polar,torch.view_as_complex, a product,torch.view_as_real). Such a model exports, but the CUDA backend cannot lower it:Inductor generates no code for complex numbers. It calls an ATen kernel for every op on a complex tensor, and this backend ships none of those kernels. Torch-TensorRT leaves
polarandview_as_realto the CUDA backend, so its programs stop at the same error.What this change does
Complex numbers as pairs of real numbers. A new pass,
ReplaceComplexWithRealPass, stores each complex tensor as a real tensor with a last dim of size 2 (the layout ofview_as_real) and rewrites the math before Inductor runs:Views, indexing, splits, concatenation and casts keep their op, with the dims and dtype of the pair. A complex constant is stored as its real view. A decomposition table cannot do this, because a decomposition must keep the dtype of the op it replaces.
The pass changes a program only when it can rewrite every complex op in it. Otherwise (a complex
add, for example) it leaves the program as it is and logs why.Complex128 becomes float64. A program that keeps float64 values in memory also needs the float64 support in #24326, so this should land after it.
Outputs that are the same tensor. When two outputs are the same tensor, AOTI puts one handle in both output slots, and the backend deleted it twice. Programs that return
y.detach(), y.detach()crashed this way. With only the complex pass, so would programs that returnview_as_realof one complex tensor twice, as Torch-TensorRT programs can. Now each repeated slot gets its own handle on the shared storage.What was tested
With torch 2.14, the GPU tests on an A100. Each new test fails without its change:
test_complex_rotary_embeddingruns a lowered rotary embedding against eager in float32 and bfloat16.test_complex_rewriteandtest_complex_rewrite_leaves_unsupported_programscheck the pass on the CPU.test_same_tensor_returned_twiceruns a program that returns one tensor twice.Rotary embeddings written as in Llama and LingBot-VA, and a Wan-style one in complex64, match eager. The Llama-style one also matches with a dynamic sequence length. The LingBot-VA ones also match through Torch-TensorRT (with torch 2.15), and Wan's complex128 one with #24326. The other
backends/cudatests that CI runs in the same pytest call still pass.