Repository navigation
[PyBindings] Make the torch-free Python bindings safe and complete for CUDA inputs and outputs - #24016
Merged
shoumikhin merged 10 commits intoOct 20, 2026
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24016
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Unclassified FailureAs of commit 11640ea with merge base d46a886 ( 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. |
shoumikhin
force-pushed
the
pybindings-followup-device-checks
branch
from
October 20, 2026 11:57
11a2109 to
046a0b7
Compare
shoumikhin
force-pushed
the
pybindings-followup-device-checks
branch
from
October 20, 2026 15:26
046a0b7 to
7479760
Compare
JacobSzwejbka
approved these changes
Oct 20, 2026
shoumikhin
force-pushed
the
pybindings-followup-device-checks
branch
from
October 20, 2026 16:57
7479760 to
7398e1b
Compare
…nsor of a run, and name input errors consistently Follow-up to the device wait rework in the Python bindings. A NumPy array or another host buffer passed for an input that the method expects on a CUDA device was accepted, and the method then read a host pointer as device memory. Torch tensors and DLPack inputs were already checked. Now a non-empty buffer is rejected the same way, before any host layout work, and the message says how to pass device memory instead. Before a run, the bindings wait in case another library still uses the memory of an input, or of an output from an earlier run. That wait covered only outputs in memory the bindings allocated, and a Module call covered no outputs at all, so a memory-planned output could be overwritten while another library was still reading it. After a run, only outputs were waited for, although a method may update an input in place. Now one helper waits, before and after each run, for every device that holds a non-empty tensor input or output of the method. Input errors now name the input and method the same way: a DLPack device mismatch names both devices; a changed input storage, in both the borrow check and the aliased torch output, says to pass the input again instead of calling set_inputs(); a wrong number or an error while choosing how to read an input gets the name; and a Method checks the input count before it reads any input. A Module call now looks its method up once. The README says that each wait is for the whole device, that it holds the GIL in CUDA-enabled bindings, how short a wait is on an idle device, that CUDA outputs come back without a copy only with MemoryPlanningPass(alloc_graph_output=False), and what a planned CUDA output costs with clone_outputs=True.
A complex32 result of these bindings exports through DLPack, but the same dtype passed back in, or from another library, was rejected as an unsupported DLPack dtype. DLPack complex with 32 bits now reads as ComplexHalf, matching the export.
ExecuTorchResult.__dlpack__ took only stream and discarded it. NumPy and other current consumers pass max_version, dl_device and copy, so calls such as np.from_dlpack(result, copy=False) failed with a TypeError, and a consumer that ordered the work itself (stream=-1) still waited for the whole device. __dlpack__ now takes the keyword-only max_version, dl_device and copy. It returns a versioned dltensor_versioned capsule when the consumer asks for DLPack 1.0, copies when asked or when the requested device differs, and raises BufferError when copy=False rules out the copy a request needs. stream=-1 skips the device wait; other streams keep the conservative wait.
A DLPack capsule, such as the one array.__dlpack__() or a result's __dlpack__() returns, was rejected as an input. Only objects with a __dlpack__ method worked. A dltensor or dltensor_versioned capsule is now accepted on its own or in the input list. Its device comes from the tensor it holds. The capsule is consumed once: it is renamed after a successful read, and passing it again raises ValueError. A read-only DLPack 1.0 CPU tensor is read from a copy, since a method may write its inputs. A producer object is now asked for a DLPack 1.0 capsule first, so it can say when its memory is read-only, with a fallback for producers that do not take max_version.
shoumikhin
force-pushed
the
pybindings-followup-device-checks
branch
from
October 20, 2026 18:34
7398e1b to
d436a11
Compare
Method.get_attribute() returned a borrowed view of any attribute, constants included. A constant loaded from a data file can sit in read-only pages, yet the result said its memory was writable, so an ordinary in-place NumPy or torch operation could crash the process. Mutable state lives in the method's planned memory, so an attribute outside it is now treated as a constant. In ExecuTorchResult output mode it is a read-only view: the buffer protocol and NumPy report it read-only, and a DLPack export is a copy, because consumers such as torch ignore the DLPack read-only flag. Only a caller that passes copy=False through DLPack 1.0 shares the memory, marked read-only. Passing the result back in as an input reads a copy. In torch output mode a constant is a copy, since a torch tensor cannot be read-only.
A Program loaded with enable_etdump=True owned its event tracer and debug
buffer, but each Method it loaded kept only a raw pointer to the tracer. So
a method from a temporary program, as in
load_program(..., enable_etdump=True).load_method("forward"), used freed
memory on every execution and could crash.
The tracer and its debug buffer now live in the program state that every
method already shares, the same way the data map does.
A torch-free caller had no way to ask whether a method expects an input on the CPU or on a CUDA device, so it had to know how the program was exported. The program already records the device of each tensor. TensorInfo gains device(), read the same way the tensor parser reads it, with tensors that carry no device annotation on the CPU. Inputs, outputs and attributes all report it. TensorInfo has grown before, with name() and is_memory_planned(); it is a value type that only MethodMeta creates. In Python, TensorInfo.device() and the new ExecuTorchResult.device give "cpu" or "cuda:N", the form ExecuTorchResult's repr already used, and TensorInfo's repr includes it. An unknown device type raises instead of being named as CUDA.
The stub for the _C extension was installed as pybindings.pyi, a name no type checker associates with _C, and the package had no py.typed marker, so the annotations never applied. The stub also imported Enum from typing, declared the Verification members without values, required inputs on Module calls that default them, and put the bundled-program methods on ExecuTorchModule, where one does not exist and the other lives on BundledModule. The wheel now ships the stub as _C.pyi, as the Buck build already names it, and a py.typed marker in executorch.extension.pybindings. The stub fixes the points above. The wheel check fails if the stub or the marker is missing.
A method whose delegate output is copied to the host, exported with MemoryPlanningPass(alloc_graph_output=False), aborted at run time, whether its inputs were on the host or on the device: the CUDA delegate got a null pointer for its output. The device propagation pass inserts a copy from each delegate output to the host and rewires the graph output to that copy, but it left the output node's spec naming the delegate output. Memory planning reads the graph outputs from that spec, so it skipped the device tensor as if it were a graph output and planned the host copy instead. The pass now updates the output node's spec to the copy.
Module.forward and the native method take clone_outputs, but the public Method wrapper of executorch.runtime always copied every output, so a caller of the recommended API could not avoid that copy. Method.execute now takes clone_outputs, with the same default and meaning as the bindings: False returns outputs that share the method's memory and may change on the next execution.
shoumikhin
force-pushed
the
pybindings-followup-device-checks
branch
from
October 20, 2026 20:36
d436a11 to
11640ea
Compare
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.
Follow-up to #23999: makes the Python bindings work well without PyTorch for programs with CUDA inputs and outputs. Each commit fixes one problem and adds a test.
What is wrong today, and the fix
MethodorModulerun now waits, before and after, for every device that holds one of its tensors. Input errors name the input.ExecuTorchResult.__dlpack__rejects the DLPack 1.0 keywords, sonp.from_dlpack(result, copy=False)fails. It now takesmax_version,dl_deviceandcopy, can return versioned capsules, and skips the device wait forstream=-1.Method.get_attribute()returns constants as writable memory, which may be read-only pages. Constants now come back read-only, DLPack exports of them are copies, and torch output mode gets a copy.Programwith ETDump on uses a freed event tracer. The tracer now lives as long as any method.TensorInfo.device()andExecuTorchResult.devicenow return"cpu"or"cuda:N"._C.pyiwithpy.typed.alloc_graph_output=False) crashes. The device copy pass now updates the graph output, so memory planning places it.executorch.runtime'sMethod.executenow takesclone_outputs, so it can skip the copy.Testing
Ran on Linux x86_64 with an H100: CUDA tests in both output modes, CPU tests with and without torch. Each new test failed on a build without its fix:
stream=-1, GPU inputs with unplanned host outputs.Method.execute(clone_outputs=False)._C.pyiorpy.typed; mypy resolves the stub from a built package.Not tested: the wait after a run for inputs updated in place.