Repository navigation
CUDA backend: lower rotary embeddings written with complex numbers - #24403
shoumikhin wants to merge 7 commits into
Conversation
🔗 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. ⏳ 5 Pending, 1 Unrelated FailureAs of commit 199a90e with merge base a057bb4 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
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. |
shoumikhin
changed the title
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 19s —— View job Nothing blocks this PR, but Detailed reviewPR Review: #24403SummaryThe PR stores complex tensors as pairs of real numbers, turns fused attention variants back into Code Quality
Testing
RecommendationApprove 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. |
|
@claude /pr-review |
|
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 reviewPR Review: #24403SummaryThe PR stores complex tensors as pairs of real numbers, turns fused attention back into Code Quality
Testing
RecommendationApprove Nothing blocks merging as long as this lands after #24326, as the description plans. If that order cannot be guaranteed, decline complex128 first. |
bbe873c to
d90c110
Compare
shoumikhin
changed the title
|
@claude /pr-review |
|
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 Detailed reviewPR Review: #24403SummaryThe PR adds Code Quality
Testing
RecommendationApprove 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. |
|
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.
|
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.
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.
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.
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.
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.
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.
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.
d90c110 to
199a90e
Compare
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 36m 48s —— View job Nothing blocks this PR, but there are three minor findings: models with several layers may still not lower through Torch-TensorRT, the new CUDA graph test never checks the output that depends on its input, and a refused device-to-device copy turns CUDA graphs off for good. Detailed reviewPR Review: #24403SummaryThe PR adds The follow-ups the author already promised are still open, and this review does not repeat them:
Code Quality
Testing
Performance
RecommendationApprove Nothing blocks merging. Settle the Torch-TensorRT question first: it decides whether the case the description cites works for whole models. |
What is wrong today
Rotary position embeddings are often written with complex numbers:
torch.polar,torch.view_as_complex, a product, thentorch.view_as_real. Such a model exports, but the CUDA backend cannot lower it:Inductor generates no code for complex numbers, so it calls ATen kernels, which this backend does not ship. 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 ofview_as_real) and rewrites the ops on it as real math:polarbecomes cosine and sine times the magnitude; a complex product becomes four real products.view_as_complexandview_as_realgo away._reshape_copy, which Torch-TensorRT makes of every view, such as Llama'sfreqs_cis.view(...).If the pass cannot rewrite every op on a complex tensor, for example a complex add, a complex input or output of the partition, a
torch.condbranch, or complex32, it leaves the program as it is and logs why. So the complex math has to stay inside one CUDA partition. Complex128 becomes float64, which the CUDA runtime supports.Also in the CUDA backend:
.item()does, and the call runs again without a graph.test_varlen_attncompared no case. It now skips only windows over query and key sequences that differ, which AOTriton lines up differently.What was tested
With torch 2.14 on an A100:
test_complex_rotary_embedding(GPU) andtest_replace_complex_with_real.py(CPU) compare rewritten programs with eager, in float32, bfloat16 and complex128, with a dynamic length,_reshape_copyand a buffer with two placeholders. Unsupported programs stay as they are.test_constant_copy_runs_without_a_graphandtest_aoti_torch_copy_cover the copy and clone refusal.The new tests fail without their fixes, and the other CUDA tests pass. The ROCm change was not run on ROCm.