Repository navigation
[PyBindings] Wait for the device instead of per-result flags, and fix bfloat16 result layout and input errors - #23999
Open
shoumikhin wants to merge 1 commit into
Open
shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
shoumikhin
requested review from
JacobSzwejbka and
larryliu0820
as code owners
October 19, 2026 18:20
shoumikhin
added
the
release notes: api
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23999
Note: Links to docs will display an error until the docs builds have been completed. ⏳ 3 Pending, 2 Unrelated FailuresAs of commit 9868f97 with merge base b1beef3 ( BROKEN TRUNK - The following jobs 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
JacobSzwejbka
approved these changes
Oct 19, 2026
… bfloat16 result layout and input errors A CUDA result borrowed from a method, through get_outputs(clone_outputs=False) or a method that returns its own input, shares the method's memory. Whether another library might still be writing that memory was tracked with a flag on each result object, and every borrowed result started with a fresh flag. So when another library wrote the memory through one result, and a second view of the same memory was passed back in, the call could read the data from before that write. The flags are gone. Before a call runs, it now waits once for each device that holds one of its CUDA inputs, including results of these bindings. A method also waits for the device of each output buffer it writes again, because another library may still be reading a dropped result that used it. After the call runs, it waits once for each device that holds an output, as before. Handing a CUDA result to another library through __dlpack__ or to_host() always waits too. In bindings built without CUDA, __dlpack__ now waits through torch, as calls already do, instead of failing. A result of these bindings passed back in now always takes the DLPack path, which keeps its real dtype, shape and strides. A bfloat16 result used to take the buffer path, whose flat bytes were read in the method's layout, so a result in another dim order came back permuted with no error. A bfloat16 result now also keeps its own shape instead of taking the method's. A failed device allocation no longer turns off the no-copy output return for the rest of a method's life, because the new buffer is allocated into a local first. A clone_outputs=False read is now its own borrowed view, not the object a default read handed over. An error raised while reading an input keeps its Python type. A Python exception gets the input and method as a note instead of being rebuilt, which failed for types such as UnicodeDecodeError. C++ runtime errors, the device checks and the unsupported-type error name the input and method in their message, and std::bad_alloc still becomes MemoryError.
shoumikhin
force-pushed
the
pybindings-cuda-wait-followup
branch
from
October 20, 2026 06:17
2e4bd22 to
9868f97
Compare
shoumikhin
changed the title
This branch has not been deployed
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 the Python bindings change that added torch-free execution and device memory support.
What was wrong
A CUDA result borrowed from a method, through
get_outputs(clone_outputs=False)or a method that returns its own input, shares the method's memory. The bindings tracked "might another library still be writing this?" with a flag on each result object, and every borrowed result started with a fresh flag. So when another library wrote the memory through one result, and a second view of the same memory was passed back in, the call could read the old data.A bfloat16 result on the host, passed back in, was read in the method's layout instead of its own. A result in another dim order came back permuted, with no error.
What this changes
The flags are replaced by device waits. Before a call runs, it waits once per device for its CUDA inputs, including results of these bindings. A method also waits for the device of an output buffer it writes again, since another library may still be reading a dropped result there. After a call, it waits once per output device.
__dlpack__andto_host()on a CUDA result always wait. In bindings built without CUDA,__dlpack__now waits through torch, as calls already do, instead of failing.A result of these bindings passed back in always takes the DLPack path, which keeps its dtype, shape and strides. This fixes the bfloat16 layout bug, and a bfloat16 result now keeps its own shape instead of taking the method's.
A failed device allocation no longer turns off the no-copy output return. A
clone_outputs=Falseread is now its own borrowed view.An error raised while reading an input keeps its Python type. A Python exception gets the input and method as a note, and C++ errors name them in the message.
The torch-free CI job now also builds a channels-last bfloat16 program, so the layout test runs there.
Tested
On Linux x86_64 with H100 GPUs, built the bindings with CUDA, without CUDA, and as the standalone torch-free module. Ran
test_pybindings_cuda_dlpack.py,test_pybindings_no_aten.py,test_pybindings.pyandruntime/test/test_runtime.pyin both output modes. Every failure there also fails without this change: tests written for the other output mode, two quantized-op tests whose kernels are not registered here, and three subprocess tests that need the CUDA runtime library outside torch. Each new or changed test fails with its fix removed.