Keep the CUDA memory pool warm between delegates - #22312
Conversation
🔗 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. |
|
ba2b029 to
ea11a4d
Compare
ea11a4d to
aab59d3
Compare
aab59d3 to
4fef258
Compare
4fef258 to
15b1c14
Compare
15b1c14 to
b4af62b
Compare
b4af62b to
edec536
Compare
edec536 to
eb407fe
Compare
eb407fe to
2b7f1f6
Compare
2b7f1f6 to
a074889
Compare
There was a problem hiding this comment.
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;
| MemPoolState& mem_pool_state() { | ||
| static MemPoolState state; | ||
| return state; | ||
| } |
There was a problem hiding this comment.
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) {
a074889 to
102e24a
Compare
There was a problem hiding this comment.
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> |
102e24a to
44b236c
Compare
There was a problem hiding this comment.
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);
}
| 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.
44b236c to
1528408
Compare
There was a problem hiding this comment.
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.txtcompilesruntime/cuda_allocator.cppinto bothaoti_cuda_shims(lines ~235-255) andaoti_cuda_backend(lines ~344-348). That makes it possible for allocations (e.g. from AOTI model code linked againstaoti_cuda_shims.dll) to populate one pool map whileCudaBackend::destroy()callsCudaAllocator::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 oneCudaAllocatorimplementation/state per process (e.g., export allocator/release symbols fromaoti_cuda_shimsand never compile this TU into the backend, or add a C-exported forwarding entrypoint inaoti_cuda_shimsthat 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.
|
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 The comment there was wrong twice, and I have rewritten it. It claimed an unsynchronized caller gets The pool and the stream could name different devices. You are right that this is a regression: Two smaller ones taken with them. The retention comment had the threshold backwards, saying it On Windows: I cannot run MSVC, so I have not been able to confirm or rule out the duplicate pool map. Not changing here: the graph trim's reach, the test-only accessor shape, and the assorted cleanups. |
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.
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.