Carry a tensor's device when wrapping it for PyTorch - #21753
Merged
Conversation
## The problem
ExecuTorch can run a model on a GPU and keep the intermediate results
(the
"activations") in GPU memory the whole way through, so nothing is copied
back and forth
across the host boundary. A caller passes GPU tensors in and gets a GPU
tensor back.
Reading that result from Python crashed:
```
Fatal Python error: Segmentation fault
File "run.py", line 16 in <module> # module.forward(inputs)
```
The model and the GPU delegate are both fine. The fault is in the small
layer that wraps
an ExecuTorch tensor as a PyTorch tensor.
An ExecuTorch tensor records which device its data lives on. That
wrapper ignored it:
```cpp
at::from_blob(etensor.mutable_data_ptr(), sizes, strides,
at::TensorOptions(dtype)); // no device, so it defaults to CPU
```
So a tensor whose storage is GPU memory came back labelled CPU, and
anything that then
touched it on the host read GPU memory directly. The Python bindings do
exactly that:
they clone each output before returning it, and that clone is a host
memory copy. The
backtrace lands in a vectorized CPU copy kernel:
```
at::native::AVX2::direct_copy_kernel(at::TensorIteratorBase&)
at::native::copy_impl(at::Tensor&, at::Tensor const&, bool)
at::native::clone(at::Tensor const&, std::optional<c10::MemoryFormat>)
```
## The change
Take the device from the tensor being wrapped instead of assuming CPU.
The construction
uses the ATen fluent builder:
```cpp
const c10::Device device = executorch_to_torch_device(etensor.device());
at::Tensor t = at::for_blob(etensor.mutable_data_ptr(), sizes)
.strides(strides)
.options(at::TensorOptions(dtype).device(device))
.target_device(device)
.make_tensor();
```
A device the build has no mapping for is rejected rather than reported
as CPU, because
treating an unreadable pointer as host memory is what caused the crash.
Both parts of the device argument are load-bearing:
- The device must go through `TensorOptions`. The tensor's dispatch key
set is derived
from options, not from `target_device`, so passing the device only as
`target_device`
produces a `cuda:0`-labelled tensor with a CPU dispatch key.
- `target_device` must also be present. Without it, ATen infers the
device by inspecting
the data pointer, and an empty tensor has no pointer to inspect: the
runtime sets it
to null when a tensor has no elements. `tensor_parser_aten.cpp` had
already hit this
and solved it the same way, so this follows that precedent.
This is a labelling fix, not a copy. The data pointer is unchanged, so a
program exported
to keep activations on the device still hands the delegate the caller's
own buffer.
## Test plan
- Ran a GPU model exported with device-resident activations on an NVIDIA
GPU. Before this
change the process died with a segmentation fault; after it, the output
matches eager
PyTorch exactly (largest absolute difference 0).
- Checked the mechanism directly against real device memory: the old
wrapping reports
`cpu` and dies with signal 11 when cloned, the new one reports `cuda:0`,
clones and the
values match.
- Confirmed no copying is introduced. A device-activation program
contains no copy
operators, its method inputs report as not memory planned so the runtime
shares the
caller's buffer, and the caller's pointers are unchanged across the
call. A
host-activation program still shows its copy operators and still matches
eager PyTorch.
- Verified the dispatch-key point directly. Constructing the wrapper
with only
`target_device(cuda:0)` (no device on `TensorOptions`) yields a tensor
whose device
reports `cuda:0` but whose key set is `DispatchKeySet(CPU, ...)`.
Constructing with the
device on both options and `target_device` yields the correct
`DispatchKeySet(CUDA, ...)`.
- Added three unit tests and verified with a mutation harness that the
empty-tensor test
(`AliasATTensorToETensorHandlesAnEmptyTensor`) fails under both
plausible regressions:
restoring the old options-only call, and passing the device only via
options without
`target_device`. The other two tests (`DeviceMapping`,
`DeviceMappingAbortsOnUnknownType`) cover the mapping helper on its own.
Not covered: no open-source continuous integration job builds this test
file. It is
compiled only where PyTorch and ATen are available, which the C++ test
job does not
enable, so the tests here are verified locally and by the internal
build. Wiring that up
means enabling the Python bindings in that job, which is a change to
shared test
infrastructure and belongs on its own.
Left for later, all pre-existing and none of them regressions from this
change:
`alias_tensor_ptr_to_attensor` and the `type_convert` helpers in the
same directory drop
the device the same way. `make_tensor_ptr` already accepts a device, so
what is missing
is only a `torch_to_executorch_device` conversion in the opposite
direction to feed it.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21753
Note: Links to docs will display an error until the docs builds have been completed. ❌ 83 Pending, 1 Unclassified FailureAs of commit 33a9c42 with merge base bd89c17 ( 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
This PR needs a
|
Gasoonjia
approved these changes
Aug 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-lands the change from #21722, which was merged into its stacked base branch rather than into main, so the fix never reached main.
Wrapping an ExecuTorch tensor for PyTorch dropped the device, so the resulting tensor carried no device label. A caller that reads the label then either fails with "tensor does not have a device" or, for accelerator memory, reads device memory from the host and faults.
Measured on an NVIDIA H100 with a CUDA delegated model whose activations stay on the GPU: before this change the Python path faults, and the same program run from C++ succeeds, which places the defect in the bindings rather than in the runtime. A CPU model is affected the same way, because the wrapped tensor has no device of either kind.
Tests: extension/aten_util/test/aten_bridge_test.cpp covers the device mapping, the abort on an unknown device type, and an empty CUDA labelled tensor, which is the case that distinguishes a correct implementation from one that only sets the label or only sets the target device.