Repository navigation
Python runtime: enforce CUDA stream selection rules on the stream itself - #24346
Open
pytorchbot wants to merge 5 commits into
Open
pytorchbot wants to merge 5 commits into
pytorchbot wants to merge 5 commits into
Conversation
…rRT delegates (#24273) ## What is wrong today A C++ program can choose the CUDA stream (a queue of GPU work) that the CUDA and TensorRT delegates run on, for example a green-context stream that keeps work on part of the GPU. Python code cannot: each delegate uses its own stream. ## What this change does This adds `executorch.runtime.cuda`, an experimental module that does not need PyTorch: ```python from executorch.runtime import cuda s = cuda.Stream() with cuda.stream(s): outputs = first.execute(inputs) # a TensorRT delegate second.execute(outputs) # a CUDA delegate in another program s.close() ``` - `cuda.Stream()` creates a stream that `close()` destroys. `cuda.ExternalStream(stream_ptr)` wraps one someone else owns. Any object with `__cuda_stream__` works too, including `torch.cuda.Stream`. - `cuda.stream(s)` selects `s` for the current thread. Blocks end on their own thread, innermost first. - Inside a block, host copies run on `s` and a call waits only for `s`. Calls with GPU inputs or outputs still wait for the whole GPU. - Limits for now: only CUDA device 0, no default stream (handle 0), no CUDA graphs inside a block, and no `await` or `yield` inside a block. A stream known to be on another device is refused. ## Tests On Linux with an A100 GPU: - Timing events show a CUDA program's kernels land on `s` inside the block, and not outside it. - Misuse raises: a plain integer, a stream on another device, the default stream, closing a selected stream (also through another wrapper), an exit out of order, an extra exit, a copy, and a closed stream. - Two threads share one `cuda.stream(s)` object, and an exit from another thread is refused. - A process that never imports PyTorch (with `EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch`) runs two CUDA-delegated programs inside one block, and the delegates see the selected stream. (cherry picked from commit fb34f59)
…elf (#24327) ## What is wrong today `executorch.runtime.cuda` (added in #24273) lets Python code choose the CUDA stream that the CUDA and TensorRT delegates run on. Some rules were checked on what a wrapper says, not on the stream, and some docs promised too much: - A stream on GPU 1, given as `cuda.ExternalStream(ptr)` or `torch.cuda.ExternalStream(ptr)`, was accepted, then failed in the delegate. - `cudaStreamLegacy` (handle 1) got past the refusal of the legacy default stream (handle 0). - The docs said a closed stream cannot be selected, a dropped scope object ends its block, and a call inside a block waits only for its stream. None of these always holds. ## What this change does - In CUDA 13 builds, the native scope asks CUDA for the stream's device when a block starts (`cudaStreamGetDevice`) and refuses any device but 0. Every CUDA wheel built from main is CUDA 13. CUDA 12 builds keep only the wrapper check, since they can load a runtime older than 12.8, which lacks that function. On a stream that is capturing a CUDA graph, the query is skipped: CUDA refuses it there, and the refusal ends the caller's capture. - The native scope refuses handles 0 and 1, `None`, an exit from another thread, and a second entry. After refusing an exit out of order, it can be entered again once its block has ended. - `_CallerStreamScope.can_check_device` is a build-time attribute, not a CUDA call. - The docs say what is enforced, which waits a block shortens, and which steps still wait longer (outputs whose layout changes, loading CUDA weights). - New tests: a block inside `torch.cuda.graph()` leaves the graph intact, a refused native scope can be entered again, and CUDA 13 CI checks that the device query is built in. - `unittest-cuda` no longer starts for a change to `runtime/cuda.py` alone; nothing it runs imports it. ## Tests On Linux with A100 GPUs, bindings built with CUDA 13.2: - `test_pybindings_cuda_dlpack.py` passes in both output modes, with a second GPU so the device test runs. - Each new check fails without its fix: capture check removed, refused exit keeping its mark, device query compiled out, and the bindings from before these changes. - A library built against CUDA 12.8 that calls `cudaStreamGetDevice` fails to load next to the CUDA 12.6 runtime. Not run on Windows, ROCm, or a real CUDA 12 build. (cherry picked from commit c030572)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24346
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit e3e0eb8 with merge base 4db778e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
pytorchbot
had a problem deploying
to
upload-benchmark-results
October 30, 2026 19:04 — with 
GitHub Actions
Failure
GitHub Actions
Failure
Base automatically changed from
cherry-pick-24273-by-pytorch_bot_bot_
to
release/1.6
October 30, 2026 21:49
JacobSzwejbka
had a problem deploying
to
upload-benchmark-results
October 31, 2026 02:35 — with
GitHub Actions
Failure
GitHub Actions
Failure
JacobSzwejbka
deployed
to
upload-benchmark-results
October 31, 2026 06:53 — with
GitHub Actions
Active
GitHub Actions
Active
Contributor
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 26m 19s —— View job Review progress
Approve: nothing blocks this cherry-pick. There is one minor regression: the private native Detailed reviewPR Review: #24346SummaryThis cherry-picks #24327 to
Docs and tests are updated to match. There are no blockers, and the issues below are minor. Thread Safety
Testing
RecommendationApprove Nothing blocks this cherry-pick. The Thread Safety item is a minor regression, and only direct users of a private class are affected. The test items can land on main and be cherry-picked later. |
The runtime.cuda docstring and the pybindings README said that loading a method's CUDA weights waits for the whole GPU. Only weights an older export keeps in one blob do that (cudaDeviceSynchronize). Current exports copy each weight with a plain cudaMemcpy, which waits for the legacy default stream and blocking streams but not for a non-blocking stream. Measured with a 2 s kernel on a non-blocking stream: load_method took 0.01 s. Docs only.
This branch was successfully 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.
What is wrong today
executorch.runtime.cuda(added in #24273) lets Python code choose the CUDA stream that the CUDA and TensorRT delegates run on. Some rules were checked on what a wrapper says, not on the stream, and some docs promised too much:cuda.ExternalStream(ptr)ortorch.cuda.ExternalStream(ptr), was accepted, then failed in the delegate.cudaStreamLegacy(handle 1) got past the refusal of the legacy default stream (handle 0).What this change does
cudaStreamGetDevice) and refuses any device but 0. Every CUDA wheel built from main is CUDA 13. CUDA 12 builds keep only the wrapper check, since they can load a runtime older than 12.8, which lacks that function. On a stream that is capturing a CUDA graph, the query is skipped: CUDA refuses it there, and the refusal ends the caller's capture.None, an exit from another thread, and a second entry. After refusing an exit out of order, it can be entered again once its block has ended._CallerStreamScope.can_check_deviceis a build-time attribute, not a CUDA call.torch.cuda.graph()leaves the graph intact, a refused native scope can be entered again, and CUDA 13 CI checks that the device query is built in.unittest-cudano longer starts for a change toruntime/cuda.pyalone; nothing it runs imports it.Tests
On Linux with A100 GPUs, bindings built with CUDA 13.2:
test_pybindings_cuda_dlpack.pypasses in both output modes, with a second GPU so the device test runs.cudaStreamGetDevicefails to load next to the CUDA 12.6 runtime.Not run on Windows, ROCm, or a real CUDA 12 build.