From d55f0928cac4e14471a52b45e7cc1fc16bc4394c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 29 Aug 2026 14:41:04 -0700 Subject: [PATCH 1/6] Keep the CUDA memory pool warm between delegates 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. --- backends/cuda/runtime/cuda_allocator.cpp | 211 +++++++++++++++++- backends/cuda/runtime/cuda_allocator.h | 44 ++++ backends/cuda/runtime/cuda_backend.cpp | 14 ++ .../cuda/runtime/test/test_cuda_allocator.cpp | 178 +++++++++++++++ extension/cuda/runtime_api.h | 1 + 5 files changed, 445 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 4c7d6aec288..7e14414511d 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -12,6 +12,12 @@ #include #include +#if !defined(EXECUTORCH_USE_HIP) +#include +#include +#include +#endif + namespace executorch::backends::cuda { using executorch::runtime::Error; @@ -21,6 +27,101 @@ using executorch::runtime::etensor::DeviceType; namespace { +#if !defined(EXECUTORCH_USE_HIP) +// The stream ordered allocator hands physical memory back to the driver +// whenever a synchronization observes a pending free, so with the default +// release threshold of zero a pool is emptied repeatedly during one inference +// and every allocation has to map memory again, which measured three orders of +// magnitude slower on an embedded board. +// +// The delegate allocates from a pool it creates rather than the device default +// pool, because the default one is shared with every other user of the async +// allocator in this process. Raising a threshold there caps what that user's +// cache may keep, and trimming it on teardown throws their cached blocks away. +// Owning the pool means the threshold and the trim only ever affect this +// backend, with no attempt to remember and restore somebody else's setting. +constexpr uint64_t kMemPoolReleaseThreshold = UINT64_MAX; + +struct MemPoolState { + std::mutex mutex; + std::unordered_map pools; +}; + +MemPoolState& mem_pool_state() { + static MemPoolState state; + return state; +} + +// Resolves the "current device" sentinel that callers are allowed to pass. +// Returns a negative value when the device cannot be determined. +int resolve_device(DeviceIndex index) { + if (index >= 0) { + return static_cast(index); + } + int current = 0; + const cudaError_t err = cudaGetDevice(¤t); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaGetDevice failed: %s. Using pool defaults.", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return -1; + } + return current; +} + +// The pool this backend allocates from on a device, creating it on first use. +// Returns nullptr when the pool cannot be created, in which case the caller +// falls back to the device default pool and only loses speed. +cudaMemPool_t mem_pool_for(int device) { + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + const auto it = state.pools.find(device); + if (it != state.pools.end()) { + return it->second; + } + + cudaMemPoolProps props{}; + props.allocType = cudaMemAllocationTypePinned; + props.handleTypes = cudaMemHandleTypeNone; + props.location.type = cudaMemLocationTypeDevice; + props.location.id = device; + + cudaMemPool_t pool = nullptr; + cudaError_t err = cudaMemPoolCreate(&pool, &props); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolCreate failed for device %d: %s. Using the default pool.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + // Recorded so a permanent failure is not retried and re-logged on every + // allocation. + state.pools.emplace(device, nullptr); + return nullptr; + } + + uint64_t threshold = kMemPoolReleaseThreshold; + err = cudaMemPoolSetAttribute( + pool, cudaMemPoolAttrReleaseThreshold, &threshold); + if (err != cudaSuccess) { + ET_LOG( + Error, + "Setting the pool release threshold failed for device %d: %s. Keeping " + "the pool at its defaults.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + + state.pools.emplace(device, pool); + return pool; +} + +#endif // !EXECUTORCH_USE_HIP + Error copy_impl( void* dst, const void* src, @@ -303,16 +404,38 @@ Result CudaAllocator::allocate_async( DeviceIndex index, cudaStream_t stream) { void* ptr = nullptr; - cudaError_t err = cudaMallocAsync(&ptr, nbytes, stream); + cudaError_t err; + // Named for the log below, so a failure points at the call that actually ran. + const char* allocator_name = "cudaMallocAsync"; + int log_device = static_cast(index); +#if defined(EXECUTORCH_USE_HIP) + err = cudaMallocAsync(&ptr, nbytes, stream); +#else + // Allocating from this backend's own pool keeps its retained memory out of + // the device default pool, which other users of the async allocator share. + const int device = resolve_device(index); + cudaMemPool_t pool = device >= 0 ? mem_pool_for(device) : nullptr; + if (device >= 0) { + log_device = device; + } + if (pool != nullptr) { + allocator_name = "cudaMallocFromPoolAsync"; + err = cudaMallocFromPoolAsync(&ptr, nbytes, pool, stream); + } else { + err = cudaMallocAsync(&ptr, nbytes, stream); + } +#endif if (err != cudaSuccess) { ET_LOG( Error, - "cudaMallocAsync failed: %s (requested %zu bytes on device %d)", + "%s failed: %s (requested %zu bytes on device %d)", + allocator_name, cudaGetErrorString(err), nbytes, - static_cast(index)); + log_device); return Error::MemoryAllocationFailed; } + return ptr; } @@ -323,6 +446,7 @@ void CudaAllocator::deallocate_async( if (ptr == nullptr) { return; } + cudaError_t err = cudaFreeAsync(ptr, stream); if (err != cudaSuccess) { ET_LOG( @@ -334,6 +458,87 @@ void CudaAllocator::deallocate_async( } } +cudaMemPool_t CudaAllocator::pool_for_device(DeviceIndex index) { +#if defined(EXECUTORCH_USE_HIP) + (void)index; + return nullptr; +#else + const int device = resolve_device(index); + if (device < 0) { + return nullptr; + } + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + const auto it = state.pools.find(device); + return it == state.pools.end() ? nullptr : it->second; +#endif +} + +void CudaAllocator::release_cached_memory(DeviceIndex index) { +#if defined(EXECUTORCH_USE_HIP) + (void)index; +#else + std::vector> targets; + { + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + if (index >= 0) { + const auto it = state.pools.find(static_cast(index)); + if (it == state.pools.end()) { + return; + } + targets.emplace_back(it->first, it->second); + } else { + // A caller asking for everything gets every pool this backend created, + // not whichever device the calling thread happens to be current on, since + // the delegate that ran is often not on that device. + targets.assign(state.pools.begin(), state.pools.end()); + } + } + + // The pools stay in the map. Trimming empties one without invalidating it, so + // a later load reuses it rather than paying to create it again. + for (const auto& [device, pool] : targets) { + if (pool != nullptr) { + // Only frees the driver has already observed can be released, so a caller + // that has not synchronized gets less back. Synchronizing here is not an + // option: the stream this backend's frees went to is the handle's own, + // and destroy() has already destroyed it by the time this runs, so + // touching it is undefined behaviour. + const cudaError_t err = cudaMemPoolTrimTo(pool, 0); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolTrimTo failed for device %d: %s.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + } + + // Memory a method allocated while its CUDA graph was being captured belongs + // to the device graph pool, which the pool trim cannot reach, so a + // graph-enabled method would otherwise hold its footprint for the life of + // the process. Outside the branch above because graph memory is a device + // resource and exists whether or not this backend has a pool here. + // + // Device scoped, unlike everything else in this function: it releases + // unused graph memory cached by every user of the device, so another + // library in this process pays to map its own graph allocations again. + // Nothing breaks, since only unused blocks go. + const cudaError_t graph_err = cudaDeviceGraphMemTrim(device); + if (graph_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceGraphMemTrim failed for device %d: %s.", + device, + cudaGetErrorString(graph_err)); + (void)cudaGetLastError(); + } + } +#endif +} + Error CudaAllocator::memcpy_async( void* dst, const void* src, diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index b0a76a51f6d..254068b8f2c 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -68,6 +68,50 @@ class CudaAllocator final : public executorch::runtime::DeviceAllocator { executorch::runtime::etensor::DeviceIndex index, cudaStream_t stream); + /** + * Return memory this backend's device pool is holding for reuse back to the + * driver. + * + * The pool keeps freed memory so that repeated allocations do not have to map + * it again, which is what makes delegate execution cheap, so a long-lived + * process should call this once its work on the device is finished. Only + * frees the driver has already observed can be released, so a caller that has + * not synchronized simply gets less back. Allocations that are still live are + * unaffected either way. + * + * The pool belongs to this backend rather than being the device default pool, + * so the pool trim never affects memory another user of the async allocator + * is holding. The graph memory trim is the exception: it is scoped to the + * device, so it also releases unused graph memory cached by other users in + * this process. + * + * Does nothing on ROCm. HIP has equivalents for all of these calls; this + * repository's CUDA-to-HIP compatibility header does not alias them yet. + * + * @param index Device to release on, or a negative value to release every + * device this backend has allocated on. Note that means every device, not + * the current one, which is what a negative value means elsewhere in this + * class: a delegate is often torn down from a thread that is not current + * on the device it ran on, so releasing only the current device would leave + * that memory held. + */ + static void release_cached_memory( + executorch::runtime::etensor::DeviceIndex index); + + /** + * The memory pool this backend allocates from on a device, or nullptr if it + * has not allocated there or the pool could not be created. + * + * Exposed so a test can observe what the pool is holding, which is not + * visible through the device default pool. No production caller; a test + * friend would be the tidier shape and would also let the HIP alias for the + * pool type go. + * + * @param index Device to query, or a negative value for the current one. + */ + static cudaMemPool_t pool_for_device( + executorch::runtime::etensor::DeviceIndex index); + /** * Copy memory asynchronously on the given CUDA stream. * Supports H2D, D2H, and D2D based on src/dst device types. diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 4322d2da15a..4b3c5d5dcf4 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -470,6 +470,8 @@ class ET_EXPERIMENTAL CudaBackend final mutable_state_note_handle(handle); + live_handles_.fetch_add(1, std::memory_order_acq_rel); + return (DelegateHandle*)handle; // Return the handle post-processing } @@ -892,6 +894,14 @@ class ET_EXPERIMENTAL CudaBackend final } delete handle; + + // The allocator lets the device pool keep freed memory so that repeated + // delegate execution does not pay to map it again. Nothing is running on + // this backend once the last handle is gone, so hand that memory back + // rather than hold it for the life of the process. + if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + CudaAllocator::release_cached_memory(-1); + } } private: @@ -906,6 +916,10 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex cuda_stream_mutex_; std::shared_ptr shared_cuda_stream_ = nullptr; + // Delegates alive right now. The device memory pool is shared, so it can only + // be released once none of them are left. + mutable std::atomic live_handles_{0}; + // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default // OFF; versioned FQN artifacts do not consult this option. diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index a872099c4f7..b5841ae21c8 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -181,3 +182,180 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostOnMissingDeviceFails) { a.deallocate(dptr, 0); } + +// The pool attributes these exercise have no HIP equivalent in the +// compatibility header, and the allocator's pool code is compiled out on ROCm +// for the same reason, so there is nothing to test there. +#if !defined(EXECUTORCH_USE_HIP) + +namespace { +uint64_t reserved_bytes(cudaMemPool_t pool) { + uint64_t reserved = 0; + EXPECT_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &reserved), + cudaSuccess); + return reserved; +} +} // namespace + +// The delegate allocates from a pool it owns, so its retained memory must not +// land in the device default pool that other users of the async allocator +// share. +TEST_F(CudaAllocatorTest, AllocatesFromItsOwnPool) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr) << "the allocator should have created its own pool"; + cudaMemPool_t default_pool = nullptr; + ASSERT_EQ(cudaDeviceGetMemPool(&default_pool, 0), cudaSuccess); + EXPECT_NE(owned, default_pool) << "the pool must not be the device default"; + + // The live block is reserved in the owned pool, which is what identifies it + // as the pool actually serving this allocation. + EXPECT_GE(reserved_bytes(owned), kBytes); + + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Freed memory is kept so repeated allocation stays cheap, which means a plain +// free no longer shrinks the pool. Without an explicit release a long lived +// process would hold that memory after every program was gone. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryReturnsPoolMemory) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr); + // Freed and synchronized, and still held, which is the point of the change. + ASSERT_GT(reserved_bytes(owned), 0u) + << "the pool should hold the freed block for reuse"; + + CudaAllocator::release_cached_memory(0); + + EXPECT_EQ(reserved_bytes(owned), 0u) + << "released memory should go back to the driver"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Releasing must not disturb allocations that are still in use. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryKeepsLiveAllocations) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // Large enough that the two blocks land in separate driver reservations. At a + // few megabytes they share one, so nothing can be released while either is + // live. + constexpr size_t kBytes = 64u << 20; + auto live = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(live.ok()); + auto temp = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(temp.ok()); + CudaAllocator::deallocate_async(temp.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr); + const uint64_t before = reserved_bytes(owned); + + CudaAllocator::release_cached_memory(0); + + // The freed block goes back and the live one stays reserved, so the pool + // gives up only what is not in use. + const uint64_t after = reserved_bytes(owned); + EXPECT_LT(after, before) << "the freed block should have been released"; + EXPECT_GE(after, kBytes) << "the live block must still be reserved"; + + EXPECT_EQ(cudaMemsetAsync(live.get(), 0, kBytes, stream), cudaSuccess); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + CudaAllocator::deallocate_async(live.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// A negative index releases every device this backend has allocated on, not the +// current one. This runner has a single GPU, so the two cannot be told apart +// here; what it pins is that the sentinel is resolved rather than passed to the +// driver. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsTheAllDevicesSentinel) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(-1); + ASSERT_NE(owned, nullptr) << "the sentinel should resolve to this device"; + ASSERT_GT(reserved_bytes(owned), 0u) + << "the pool should hold the freed block for reuse"; + + CudaAllocator::release_cached_memory(-1); + + EXPECT_EQ(reserved_bytes(owned), 0u); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Memory allocated during a graph capture goes to the device graph pool, which +// the pool trim cannot reach, so releasing has to trim that too. Without the +// graph trim this is the only new test that fails. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryReturnsGraphMemory) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 64u << 20; + cudaGraph_t graph = nullptr; + cudaGraphExec_t graph_exec = nullptr; + ASSERT_EQ( + cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), + cudaSuccess); + auto captured = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(captured.ok()); + CudaAllocator::deallocate_async(captured.get(), 0, stream); + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + ASSERT_EQ( + cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0), + cudaSuccess); + ASSERT_EQ(cudaGraphLaunch(graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + size_t reserved = 0; + ASSERT_EQ( + cudaDeviceGetGraphMemAttribute( + 0, cudaGraphMemAttrReservedMemCurrent, &reserved), + cudaSuccess); + ASSERT_GT(reserved, 0u) << "the capture should have reserved graph memory"; + + ASSERT_EQ(cudaGraphExecDestroy(graph_exec), cudaSuccess); + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + + CudaAllocator::release_cached_memory(0); + + ASSERT_EQ( + cudaDeviceGetGraphMemAttribute( + 0, cudaGraphMemAttrReservedMemCurrent, &reserved), + cudaSuccess); + EXPECT_EQ(reserved, 0u) << "graph memory should go back to the driver"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +#endif // !EXECUTORCH_USE_HIP diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index bae5c6a79bf..af404423548 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -18,6 +18,7 @@ using cudaError_t = hipError_t; using cudaGraph_t = hipGraph_t; using cudaGraphExec_t = hipGraphExec_t; using cudaMemcpyKind = hipMemcpyKind; +using cudaMemPool_t = hipMemPool_t; using cudaMemoryType = hipMemoryType; using cudaStreamCaptureMode = hipStreamCaptureMode; using cudaStream_t = hipStream_t; From 6bae9bc6712ad5f746a965668a27b6e89b0a90ea Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 14:44:02 -0700 Subject: [PATCH 2/6] Wait before the trim, and only use a pool the stream can reach 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. --- backends/cuda/runtime/cuda_allocator.cpp | 33 +++++++++++++------ backends/cuda/runtime/cuda_backend.cpp | 15 +++++++++ .../cuda/runtime/test/test_cuda_allocator.cpp | 26 +++++++++++++++ 3 files changed, 64 insertions(+), 10 deletions(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 7e14414511d..489aa3b4893 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -36,10 +36,11 @@ namespace { // // The delegate allocates from a pool it creates rather than the device default // pool, because the default one is shared with every other user of the async -// allocator in this process. Raising a threshold there caps what that user's -// cache may keep, and trimming it on teardown throws their cached blocks away. -// Owning the pool means the threshold and the trim only ever affect this -// backend, with no attempt to remember and restore somebody else's setting. +// allocator in this process. Raising the threshold there would make that shared +// pool hold on to memory on their behalf, and trimming it on teardown would +// throw their cached blocks away. Owning the pool means the threshold and the +// trim only ever affect this backend, with no attempt to remember and restore +// somebody else's setting. constexpr uint64_t kMemPoolReleaseThreshold = UINT64_MAX; struct MemPoolState { @@ -413,8 +414,20 @@ Result CudaAllocator::allocate_async( #else // Allocating from this backend's own pool keeps its retained memory out of // the device default pool, which other users of the async allocator share. + // + // The pool has to belong to the device the stream runs on. A pool from + // another device returns a pointer that stream cannot touch, and the failure + // surfaces later as an illegal access rather than here. The caller's index + // and the stream can disagree, so the plain async allocation is used unless + // the index names the device that is current for this stream. const int device = resolve_device(index); - cudaMemPool_t pool = device >= 0 ? mem_pool_for(device) : nullptr; + int stream_device = -1; + if (cudaGetDevice(&stream_device) != cudaSuccess) { + (void)cudaGetLastError(); + stream_device = -1; + } + cudaMemPool_t pool = + (device >= 0 && device == stream_device) ? mem_pool_for(device) : nullptr; if (device >= 0) { log_device = device; } @@ -500,11 +513,11 @@ void CudaAllocator::release_cached_memory(DeviceIndex index) { // a later load reuses it rather than paying to create it again. for (const auto& [device, pool] : targets) { if (pool != nullptr) { - // Only frees the driver has already observed can be released, so a caller - // that has not synchronized gets less back. Synchronizing here is not an - // option: the stream this backend's frees went to is the handle's own, - // and destroy() has already destroyed it by the time this runs, so - // touching it is undefined behaviour. + // Only frees the driver has already observed can be released, and a free + // still pending gives back nothing rather than less. The backend waits on + // its stream during teardown for that reason, before dropping it. This + // cannot wait on anything itself: it holds no stream, and a caller + // reaching it directly is responsible for having synchronized. const cudaError_t err = cudaMemPoolTrimTo(pool, 0); if (err != cudaSuccess) { ET_LOG( diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 4b3c5d5dcf4..1d40bf17bbd 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -866,6 +866,21 @@ class ET_EXPERIMENTAL CudaBackend final mutable_state_forget_handle(handle); + // Waited on before the stream reference goes, so the frees this handle + // queued are observed by the driver. Without this the pool trim at the end + // of teardown sees them as still pending and gives back nothing, and the + // work is being abandoned anyway. + if (handle->cuda_stream != nullptr && *handle->cuda_stream != nullptr) { + const cudaError_t sync_err = cudaStreamSynchronize(*handle->cuda_stream); + if (sync_err != cudaSuccess) { + ET_LOG( + Error, + "cudaStreamSynchronize failed during teardown: %s.", + cudaGetErrorString(sync_err)); + (void)cudaGetLastError(); + } + } + // The CUDA stream is managed by shared_ptr in the handle. // It will be automatically destroyed when the last handle using it // is destroyed. Just reset our reference. diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index b5841ae21c8..9ea5e392884 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -202,6 +202,32 @@ uint64_t reserved_bytes(cudaMemPool_t pool) { // The delegate allocates from a pool it owns, so its retained memory must not // land in the device default pool that other users of the async allocator // share. +// The retention threshold is the whole point of owning a pool: at the default +// of zero the driver empties it on every synchronize. Nothing else in this +// suite notices a smaller value, so it is asserted directly. +TEST_F(CudaAllocatorTest, PoolRetainsMemoryWithoutLimit) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + auto res = CudaAllocator::allocate_async(8u << 20, 0, stream); + ASSERT_TRUE(res.ok()); + + cudaMemPool_t pool = CudaAllocator::pool_for_device(0); + ASSERT_NE(pool, nullptr); + + uint64_t threshold = 0; + ASSERT_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReleaseThreshold, &threshold), + cudaSuccess); + EXPECT_EQ(threshold, UINT64_MAX) + << "the pool must hold on to freed memory rather than return it"; + + ASSERT_EQ(CudaAllocator::deallocate_async(res.get(), 0, stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + TEST_F(CudaAllocatorTest, AllocatesFromItsOwnPool) { cudaStream_t stream; ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); From 6d03ba2deb88412b21526fbb1e0ca4bf0e1b9c18 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 14:50:45 -0700 Subject: [PATCH 3/6] Drop an unused include and stop the pool leaking between tests 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. --- backends/cuda/runtime/test/test_cuda_allocator.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index 9ea5e392884..73adb16831e 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -8,7 +8,6 @@ #include -#include #include #include @@ -35,6 +34,15 @@ class CudaAllocatorTest : public testing::Test { } } + // The pool is meant to stay warm, so without this a test that measured + // reserved bytes would see whatever an earlier one left behind, and the order + // would matter. + void TearDown() override { + if (device_count_ > 0) { + CudaAllocator::release_cached_memory(-1); + } + } + // One past the last valid device ordinal, so switching to it always fails. // Only the tests that need such an ordinal call this, so the fit check lives // here rather than in SetUp, where it would also skip the device-0 tests. From 7c6798e70c57d7e0bf256abf30da18b29a9144c2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 15:40:21 -0700 Subject: [PATCH 4/6] Drop the HIP alias for the memory pool type The pool type accessor is only there for a test, and every test that calls it is already compiled out on ROCm, so the alias existed to satisfy a declaration nothing on that build could reach. Guarding the declaration and the definition lets the alias go, and removes the HIP branch of the accessor that returned nullptr for a caller that never existed there. Test plan: Preprocessed the allocator header, the allocator source and the allocator test both ways. With EXECUTORCH_USE_HIP defined, the pool type appears zero times in all three; without it, six, one and seven times. So the type is genuinely unreachable on ROCm and still required on CUDA. Compiled the guarded shape both ways as well: the CUDA build keeps the accessor, and a build with the alias absent and the guard active compiles without it. clang-format is clean on all three files. --- backends/cuda/runtime/cuda_allocator.cpp | 7 ++----- backends/cuda/runtime/cuda_allocator.h | 9 ++++++--- extension/cuda/runtime_api.h | 1 - 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 489aa3b4893..6d046764953 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -471,11 +471,8 @@ void CudaAllocator::deallocate_async( } } +#if !defined(EXECUTORCH_USE_HIP) cudaMemPool_t CudaAllocator::pool_for_device(DeviceIndex index) { -#if defined(EXECUTORCH_USE_HIP) - (void)index; - return nullptr; -#else const int device = resolve_device(index); if (device < 0) { return nullptr; @@ -484,8 +481,8 @@ cudaMemPool_t CudaAllocator::pool_for_device(DeviceIndex index) { const std::lock_guard lock(state.mutex); const auto it = state.pools.find(device); return it == state.pools.end() ? nullptr : it->second; -#endif } +#endif // !EXECUTORCH_USE_HIP void CudaAllocator::release_cached_memory(DeviceIndex index) { #if defined(EXECUTORCH_USE_HIP) diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index 254068b8f2c..8bf383fcb51 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -98,19 +98,22 @@ class CudaAllocator final : public executorch::runtime::DeviceAllocator { static void release_cached_memory( executorch::runtime::etensor::DeviceIndex index); +#if !defined(EXECUTORCH_USE_HIP) /** * The memory pool this backend allocates from on a device, or nullptr if it * has not allocated there or the pool could not be created. * * Exposed so a test can observe what the pool is holding, which is not - * visible through the device default pool. No production caller; a test - * friend would be the tidier shape and would also let the HIP alias for the - * pool type go. + * visible through the device default pool. No production caller. + * + * Not declared on ROCm: the pool code is compiled out there, so there is + * nothing to observe and the pool type needs no HIP alias. * * @param index Device to query, or a negative value for the current one. */ static cudaMemPool_t pool_for_device( executorch::runtime::etensor::DeviceIndex index); +#endif // !EXECUTORCH_USE_HIP /** * Copy memory asynchronously on the given CUDA stream. diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index af404423548..bae5c6a79bf 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -18,7 +18,6 @@ using cudaError_t = hipError_t; using cudaGraph_t = hipGraph_t; using cudaGraphExec_t = hipGraphExec_t; using cudaMemcpyKind = hipMemcpyKind; -using cudaMemPool_t = hipMemPool_t; using cudaMemoryType = hipMemoryType; using cudaStreamCaptureMode = hipStreamCaptureMode; using cudaStream_t = hipStream_t; From 740b171c7d477653f9ac0595a77a9ce54d425cdb Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 17:02:41 -0700 Subject: [PATCH 5/6] Call deallocate_async as a statement, since it returns void The new threshold test wrapped it in ASSERT_EQ, which cannot compile: the function returns void and the macro needs a comparable value. Every other call in the file already treats it as a statement. Reproduced with a small clang case: the ASSERT_EQ form gives two errors, the statement form none. --- backends/cuda/runtime/test/test_cuda_allocator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index 73adb16831e..4ffa81a98e0 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -231,7 +231,7 @@ TEST_F(CudaAllocatorTest, PoolRetainsMemoryWithoutLimit) { EXPECT_EQ(threshold, UINT64_MAX) << "the pool must hold on to freed memory rather than return it"; - ASSERT_EQ(CudaAllocator::deallocate_async(res.get(), 0, stream), Error::Ok); + CudaAllocator::deallocate_async(res.get(), 0, stream); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } From 109b84eaa47f0e4ee3b1e0fc750d8b7bc101cafc Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 17:21:22 -0700 Subject: [PATCH 6/6] Wait for the device, not a stream, before releasing the pool The wait this change added at the top of teardown was aimed at the wrong stream, and not only because of the change that moves methods onto the per-thread stream. The frees the trim needs the driver to have observed are issued by the storage free path, which reads the current stream at free time rather than any stream recorded on the handle. So on a caller supplied stream those frees never went to the stream being waited on, and the trim gave back nothing. The handle's stream was right by luck on one path of three. The wait now sits inside the branch that releases the pool, immediately before it, as a device synchronize. That covers whichever stream a free actually went to, matches the existing shape elsewhere in this backend where a device synchronize precedes a free for the same reason, and costs nothing on the common path because it only runs when the last delegate is going away. The live handle counter moves next to the other process wide flag rather than sitting among the shared stream members. Both are small, and together they mean this file no longer conflicts with the per-thread stream change in either direction: the two merge to a byte identical result whichever lands first, so nobody has to resolve a teardown by hand and accidentally drop the wait. --- backends/cuda/runtime/cuda_backend.cpp | 37 +++++++++++++------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 1d40bf17bbd..f826633984e 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -866,21 +866,6 @@ class ET_EXPERIMENTAL CudaBackend final mutable_state_forget_handle(handle); - // Waited on before the stream reference goes, so the frees this handle - // queued are observed by the driver. Without this the pool trim at the end - // of teardown sees them as still pending and gives back nothing, and the - // work is being abandoned anyway. - if (handle->cuda_stream != nullptr && *handle->cuda_stream != nullptr) { - const cudaError_t sync_err = cudaStreamSynchronize(*handle->cuda_stream); - if (sync_err != cudaSuccess) { - ET_LOG( - Error, - "cudaStreamSynchronize failed during teardown: %s.", - cudaGetErrorString(sync_err)); - (void)cudaGetLastError(); - } - } - // The CUDA stream is managed by shared_ptr in the handle. // It will be automatically destroyed when the last handle using it // is destroyed. Just reset our reference. @@ -915,6 +900,20 @@ class ET_EXPERIMENTAL CudaBackend final // this backend once the last handle is gone, so hand that memory back // rather than hold it for the life of the process. if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + // Only frees the driver has already observed can be released, so without + // this the trim gives back nothing rather than less. A device wait rather + // than a stream wait because the frees went to whichever stream was + // current when the storage was released, which is not necessarily one + // this handle recorded, and teardown is not bound to the thread that ran + // the method either. + const cudaError_t sync_err = cudaDeviceSynchronize(); + if (sync_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceSynchronize before releasing the pool failed: %s.", + cudaGetErrorString(sync_err)); + (void)cudaGetLastError(); + } CudaAllocator::release_cached_memory(-1); } } @@ -931,15 +930,15 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex cuda_stream_mutex_; std::shared_ptr shared_cuda_stream_ = nullptr; - // Delegates alive right now. The device memory pool is shared, so it can only - // be released once none of them are left. - mutable std::atomic live_handles_{0}; - // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default // OFF; versioned FQN artifacts do not consult this option. std::atomic weight_sharing_across_methods_{false}; + // Delegates alive right now. The device memory pool is shared, so it can only + // be released once none of them are left. + mutable std::atomic live_handles_{0}; + // --------------------------------------------------------------- // Per-weight constant cache. //