Keep the CUDA memory pool warm between delegates

[shoumikhin]

The CUDA delegate allocates through the stream ordered allocator, whose pool hands
physical memory back to the driver whenever a synchronization observes a pending
free. With the default release threshold of zero that happens repeatedly during one
inference, so nearly every allocation has to map memory again.

The delegate now allocates from a pool it creates rather than the device default
pool, with the threshold set so the pool keeps what it has. Owning the pool is what
makes that safe: the default pool is shared with every other user of the async
allocator in the process, so raising the threshold there would make that pool retain
memory the other user expected to get back, and trimming it on teardown would throw
their cached blocks away. A pool of its own means the threshold and the trim only
affect this backend, and there is nothing to remember or restore.

Because the memory is then held rather than returned at each synchronize, the backend
gives it back explicitly when the last delegate handle is destroyed. Only frees the
driver has already observed can be released, so a caller that has not synchronized
gets less back rather than anything worse, which is why this does not synchronize the
device itself: that would wait on every stream on the device, including work this
backend never queued.

Memory a method allocated while its CUDA graph was being captured belongs to the
device graph pool, which a pool trim cannot reach, so the release trims that too or a
graph enabled method would hold its footprint for the life of the process. That trim
is the one part of the release that is not isolated: it is scoped to the device, so it
also releases unused graph memory cached by other users in this process, who then pay
to map it again. The header says so at the call it applies to.

Test plan:

Five tests in backends/cuda/runtime/test/test_cuda_allocator.cpp, and the point of
each is a mutation that kills it:

pool serves allocations, not the default pool forcing pool creation to fail
a freed block is still reserved after a sync dropping the release threshold
releasing returns it dropping the pool trim
a release leaves a live allocation reserved dropping the pool trim
graph memory goes back after a release dropping the graph trim

All fourteen tests in that file pass against the change. Deleting only the graph trim
fails only the graph test, and each of the other three mutations above fails at least
three of the five, so no single one of them is carrying the suite.

Measured, per allocation, allocating and freeing with a synchronize between:

Orin Nano 3790.79 us before, 2.34 us after
Thor 363.00 us before, 1.51 us after
H100 40.30 us before, 1.01 us after
A100 18.60 us before, 1.03 us after

A private pool measured the same warm allocation cost as the default one, 1.28 us
against 1.31 us on an H100, and trimming it left a co-tenant's 256 MiB cache in the
default pool untouched. A model split into 25 delegates went from about 714 to about
518 microseconds median on an H100.

Retaining the pool means a long lived process holds that memory until its last
delegate goes away, which is visible to other processes on the same GPU. A server that
keeps a model loaded never reaches that point.

The pool calls are compiled out on ROCm and the change is a no-op there. HIP has
equivalents for all of them; this repository's compatibility header does not alias
them yet, which is the only reason for the guards.

Not covered, and worth knowing before this lands:

The release only returns blocks whose frees the driver has already observed, and
nothing on the teardown path waits for the frees this backend queued, so in the common
configurations it gives back less than the whole pool. Synchronizing there is not
available: those frees go to the handle's own stream, which destroy() has already
destroyed by the time the release runs, so touching it segfaults. Making this reliable
means freeing on a stream this backend still owns at that point, which is a change to
teardown rather than to the allocator.

The all-devices meaning of a negative index is exercised on a one-GPU runner, where it
cannot be told apart from current-device-only. The backend counter that decides when to
release, and the release call site itself, are not covered by any test in this
directory, since nothing here builds the backend.

Windows: the build compiles this file into both the shims library and the backend on
MSVC, and the pool map is a function-local static, so that build plausibly gets two
maps with the allocations in one and the trim in the other, which would make the
release a no-op there rather than merely wasteful. Both Windows CUDA jobs are skipped
for pull requests from a fork, so nothing here has exercised it and it needs someone
with that toolchain.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22312

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 2 Unrelated Failures

As of commit 102e24a with merge base c27baa8 ():

NEW FAILURES - The following jobs have failed:

• pull / android / run-emulator (gh)
  adb: failed to install android-test-debug-androidTest.apk: cmd: Can't find service: package
• pull / unittest / linux / linux-job (gh)
  pytorch_tokenizers
• pull / unittest-buck / linux / linux-job (gh)
  pytorch_tokenizers
• pull / unittest-editable / linux / linux-job (gh)
  pytorch_tokenizers

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-arm-backend-no-driver (test_pytest_models_tosa) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[linux-foundation-easycla]

• ❌ The email address for the commit (102e24a) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

backends/cuda/runtime/cuda_allocator.cpp:28

• Including <executorch/backends/aoti/slim/cuda/guard.h> brings a DeviceIndex alias into executorch::backends::cuda, and this TU also declares using executorch::runtime::etensor::DeviceIndex; in the same namespace. That results in a duplicate type-alias declaration (and can fail to compile depending on toolchain).

using executorch::runtime::Error;
using executorch::runtime::Result;
using executorch::runtime::etensor::DeviceIndex;
using executorch::runtime::etensor::DeviceType;

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (2)

backends/cuda/runtime/cuda_allocator.cpp:55

• mem_pool_state() is a function-local static, so if this translation unit is compiled into multiple DLLs (as happens on MSVC where runtime/cuda_allocator.cpp is built into both aoti_cuda_shims and aoti_cuda_backend), each DLL will get its own pool map. That can make release_cached_memory() a no-op for allocations done via the other copy, leaving the pool warm indefinitely in some Windows configurations.

MemPoolState& mem_pool_state() {
  static MemPoolState state;
  return state;
}

backends/cuda/runtime/test/test_cuda_allocator.cpp:293

• This comment/test name says a negative index means “current device”, but release_cached_memory() treats any negative index as “release all devices this backend has allocated on” (and only pool_for_device(-1) resolves to the current device). This is misleading and makes the test’s intent unclear.

// A negative index means whichever device is current, which the allocator
// resolves rather than passing on to the driver.
TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsCurrentDeviceSentinel) {

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

backends/cuda/runtime/test/test_cuda_allocator.cpp:294

• The comment/test name says a negative index means "current device", but CudaAllocator::release_cached_memory documents negative as "all devices this backend has allocated on" (and the implementation uses negative to target all pools). This is misleading for future readers; on a single-GPU runner they happen to behave the same.

// A negative index means whichever device is current, which the allocator
// resolves rather than passing on to the driver.
TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsCurrentDeviceSentinel) {
  cudaStream_t stream;