Borrow the device through the shared CUDA guard - #4729
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
force-pushed
the
adopt-shared-cuda-guard
branch
from
September 22, 2026 19:18
62eb92d to
196f821
Compare
The ExecuTorch delegate selects a CUDA device in three places: when a program is loaded, when it runs, and when its engine is freed. Two of those three selected the device and never put the caller's back, so loading or freeing a program moved the calling thread to another card and left it there. The next unrelated allocation on that thread then landed on the wrong card. Only the third one, the call that does the work, restored anything, and it did so through a small struct written for that one function. This matters for a program split between this delegate and the CUDA one, with engines on different cards, because the thread that loads one delegate is not doing that delegate's work. ExecuTorch now ships a scope guard for this in its CUDA extension, beside the caller-stream helper this file already uses. It selects a device, restores the previous one when it goes out of scope, refuses an index it cannot select, and keys the restore on what the caller asked for rather than on whether the selection reported success, because a failed selection can still move the calling thread. All three places use it now, and the local struct is gone. The destructor deliberately ignores the result. It can run while an arena is torn down, on a thread that was working on another card, and a destructor has no way to report a failure. The guard logs one instead. The vendored build of the extension carried only the caller-stream source, so the guard's definitions would have been absent from the shared library it produces. It builds both sources now, and exports both headers. A test covers the fix. It exports a coalesced program with its engine on the second card, puts the caller on the first, and reads the current device after the load, after the method is prepared, after the run, and after the free. It skips with a reason on a machine with one card, because with one card there is nothing to switch to and so nothing to restore. Test plan: built an ExecuTorch wheel carrying the guard, installed it into a fresh environment, and built this file against that wheel alone. It compiles with no warnings, and the two symbols it needs for the guard are both exported by that wheel's own library, so a build with only the wheel can link it.
shoumikhin
force-pushed
the
adopt-shared-cuda-guard
branch
from
September 23, 2026 15:12
196f821 to
ec28845
Compare
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
2 participants
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.
The ExecuTorch delegate picks a CUDA device in three places: when a program is loaded, when it
runs, and when its engine is freed. Two of those three picked a device and never put the caller's
back.
That matters for a program split between this delegate and the CUDA one, with engines on different
cards. The thread that loads one delegate is not doing that delegate's work, so moving it and
leaving it there means the next unrelated allocation on that thread lands on the wrong card.
Neither of the leaking places is the call that does the work, which is what made this easy to
miss. Only the third one put the device back, using a small helper written for that one function.
ExecuTorch now ships a scope guard for exactly this, in its CUDA extension, next to the
caller-stream helper this file already uses. It picks a device, puts the previous one back when it
goes out of scope, turns down an index it cannot select, and decides whether to restore based on
what the caller asked for rather than on whether the selection reported success, because a failed
selection can still move the calling thread.
All three places use it now, and the local helper is gone.
The one in the destructor ignores the result on purpose. It can run while a memory arena is being
torn down, on a thread that was working on another card, and a destructor has no way to report a
failure. The guard writes a log line instead.
Testing
A new test covers the bug. It exports a coalesced program with its engine on the second card, puts
the caller on the first, and reads the current device after the load, after the method is prepared,
after the run, and after the free. On a machine with one card it skips and says why, because with
one card there is nothing to switch to and so nothing to restore.
I also built an ExecuTorch wheel carrying the guard, installed it into a fresh environment, and
built this file against that wheel alone. It compiles with no warnings, and the two symbols it
needs for the guard are both exported by that wheel's own library, so a build with only the wheel
can link it.
Before merging
This needs an ExecuTorch release that carries the guard, and the pin here moved up to it. Neither
has happened yet, so please hold this until both do.