Skip to content

Keep the CUDA memory pool warm between delegates - #22312

Open
shoumikhin wants to merge 3 commits into
pytorch:mainfrom
shoumikhin:fix/cuda-mempool-release-threshold
Open

Keep the CUDA memory pool warm between delegates#22312
shoumikhin wants to merge 3 commits into
pytorch:mainfrom
shoumikhin:fix/cuda-mempool-release-threshold

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

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
Jetson AGX 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.

@shoumikhin shoumikhin added the release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs label Aug 29, 2026
Copilot AI lite review requested due to automatic review settings August 29, 2026 17:52
@shoumikhin shoumikhin added the release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs label Aug 29, 2026
@pytorch-bot

pytorch-bot Bot commented Aug 29, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 29, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@linux-foundation-easycla

linux-foundation-easycla Bot commented Aug 29, 2026

Copy link
Copy Markdown

CLA Missing ID

  • ✅ login: shoumikhin / name: Anthony Shoumikhin (0211c01)
  • ❌ The email address for the commit (1528408) 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.

@shoumikhin
shoumikhin marked this pull request as draft August 29, 2026 19:36
@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from ba2b029 to ea11a4d Compare August 29, 2026 21:41
@shoumikhin shoumikhin changed the title Keep the CUDA memory pool from releasing memory between delegates Keep the CUDA memory pool warm between delegates Aug 29, 2026
@shoumikhin
shoumikhin marked this pull request as ready for review August 29, 2026 21:41
Copilot AI review requested due to automatic review settings August 29, 2026 21:41

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from ea11a4d to aab59d3 Compare August 30, 2026 18:02
Copilot AI review requested due to automatic review settings August 30, 2026 18:02

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from aab59d3 to 4fef258 Compare August 30, 2026 19:27
Copilot AI review requested due to automatic review settings August 30, 2026 19:27

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from 4fef258 to 15b1c14 Compare August 31, 2026 05:10
Copilot AI review requested due to automatic review settings August 31, 2026 05:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from 15b1c14 to b4af62b Compare August 31, 2026 15:22
Copilot AI review requested due to automatic review settings August 31, 2026 15:22

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from b4af62b to edec536 Compare August 31, 2026 16:02
Copilot AI review requested due to automatic review settings August 31, 2026 16:02

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from edec536 to eb407fe Compare August 31, 2026 22:41
Copilot AI review requested due to automatic review settings August 31, 2026 22:41

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from eb407fe to 2b7f1f6 Compare September 1, 2026 00:00
Copilot AI review requested due to automatic review settings September 1, 2026 00:00
@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from 2b7f1f6 to a074889 Compare September 1, 2026 00:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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;

Comment on lines +52 to +55
MemPoolState& mem_pool_state() {
static MemPoolState state;
return state;
}
Comment thread backends/cuda/runtime/test/test_cuda_allocator.cpp Outdated
Copilot AI review requested due to automatic review settings September 1, 2026 00:07

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 AI review requested due to automatic review settings September 1, 2026 00:57
@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from a074889 to 102e24a Compare September 1, 2026 00:57

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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;


#include <gtest/gtest.h>

#include <executorch/backends/aoti/slim/cuda/guard.h>
Copilot AI review requested due to automatic review settings September 1, 2026 17:00
@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from 102e24a to 44b236c Compare September 1, 2026 17:00

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Suppressed comments (2)

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

  • This test name/comment describe negative indices as meaning “current device”, but release_cached_memory() documents negative indices as “release every device this backend has allocated on” (the opposite semantic). On single-GPU runners the behavior is indistinguishable, so the current wording is misleading.
// A negative index means whichever device is current, which the allocator
// resolves rather than passing on to the driver.
TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsCurrentDeviceSentinel) {

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

  • After this test completes, the allocator's private pool can still retain freed memory (by design). Trimming it at the end avoids cross-test interference when tests are shuffled or when new tests are added later.
  CudaAllocator::deallocate_async(live.get(), 0, stream);
  ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess);
  ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess);
}

Comment on lines +223 to +226
CudaAllocator::deallocate_async(res.get(), 0, stream);
ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess);
ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess);
}
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.
Copilot AI review requested due to automatic review settings September 1, 2026 17:14
@shoumikhin
shoumikhin force-pushed the fix/cuda-mempool-release-threshold branch from 44b236c to 1528408 Compare September 1, 2026 17:14

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Suppressed comments (1)

backends/cuda/runtime/cuda_allocator.cpp:53

  • mem_pool_state() uses a function-local static, so every binary/shared-library that compiles this TU gets its own pool map. On MSVC, backends/cuda/CMakeLists.txt compiles runtime/cuda_allocator.cpp into both aoti_cuda_shims (lines ~235-255) and aoti_cuda_backend (lines ~344-348). That makes it possible for allocations (e.g. from AOTI model code linked against aoti_cuda_shims.dll) to populate one pool map while CudaBackend::destroy() calls CudaAllocator::release_cached_memory(-1) in the backend’s copy, trimming a different set of pools and effectively turning the “release on last handle” path into a no-op (or trimming only part of the memory). Consider ensuring there is exactly one CudaAllocator implementation/state per process (e.g., export allocator/release symbols from aoti_cuda_shims and never compile this TU into the backend, or add a C-exported forwarding entrypoint in aoti_cuda_shims that the backend calls).
struct MemPoolState {
  std::mutex mutex;
  std::unordered_map<int, cudaMemPool_t> pools;
};

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

Review found the release giving back nothing in the case it exists for, measured on real hardware:
with a free still pending the trim recovered 0 of 256 MiB, three runs out of three, and one stream
synchronize recovered all of it. Teardown now waits on the handle's stream before dropping the
reference, which is correct on its own terms since that work is being abandoned anyway.

The comment there was wrong twice. It said an unsynchronized caller gets less back, when it gets
nothing, and it gave "the stream is already destroyed" as the reason a wait is impossible. In shared
stream mode the backend never resets its own reference, so a stream was alive the whole time.

Second, the pool is chosen from the caller's device index while the stream still orders the work, and
nothing checked the two name the same device. A pool from another device returns a pointer that
stream cannot touch, and it surfaces later as an illegal access rather than at the allocation. The
backend can produce that mismatch on its own, since the execution stream is filed under a fixed key.
The private pool is now used only when the index names the current device, and the plain async
allocation covers the rest, which is what happened before this change.

Also corrected the retention comment, which said raising the threshold caps what a cache may keep.
It does the opposite: the driver holds that many bytes before releasing to the OS.

Test plan:

    threshold pinned at UINT64_MAX      new test, the suite passed at any value before
    clang-format                        clean on all three files

The threshold test is the one gap worth naming: nothing else in the suite notices a smaller value,
so the change's whole purpose was unpinned.
Copilot AI review requested due to automatic review settings September 1, 2026 21:44

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Thanks, the first two were real and I have fixed them.

The release gave back nothing, not less. Reproduced: with a free still pending the trim recovered
0 of 256 MiB, and one stream synchronize recovered all of it. Teardown now waits on the handle's
stream before dropping the reference. That is correct on its own terms too, since the work is being
abandoned at that point.

The comment there was wrong twice, and I have rewritten it. It claimed an unsynchronized caller gets
less back, and it gave "the stream is already destroyed" as the reason a wait was impossible. In
shared stream mode the backend never resets its own reference, so a stream was alive the whole time.
The description said the same thing and is corrected.

The pool and the stream could name different devices. You are right that this is a regression:
the old call always allocated from the stream's own device, and the new one takes the pool from the
caller's index with nothing checking they agree. The private pool is now used only when the index
names the current device, and the plain async allocation covers the rest, which restores the previous
behaviour rather than papering over the mismatch.

Two smaller ones taken with them. The retention comment had the threshold backwards, saying it
caps what a cache may keep when it does the opposite. And nothing pinned the threshold, so a much
smaller value passed the whole suite. There is a test for it now.

On Windows: I cannot run MSVC, so I have not been able to confirm or rule out the duplicate pool map.
If someone with that toolchain can check it, that is the one open item I would want a second pair of
eyes on before this lands.

Not changing here: the graph trim's reach, the test-only accessor shape, and the assorted cleanups.
Each is worth its own change rather than growing this one.

The guard header was included but nothing in the file uses any of its four functions.
CallerStreamGuard, the one guard type the test does use, comes from caller_stream.h. Removing it also
removes a header the test target does not declare a dependency on.

And the fixture had no teardown, so a pool left warm by one test was still warm for the next. That is
the intended production behaviour but it makes a test that measures reserved bytes depend on order,
which matters under shuffle. It releases after each test now.
Copilot AI review requested due to automatic review settings September 1, 2026 21:50

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants