From 15b1c14f2a54f1ea5c8598b6dd6fc8bcdc90cd14 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 29 Aug 2026 14:41:04 -0700 Subject: [PATCH] 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 a threshold there caps what their cache may keep, and trimming it on teardown throws their cached blocks away. A pool of its own means the threshold and the trim only ever 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. Test plan: Four tests in backends/cuda/runtime/test/test_cuda_allocator.cpp. They check that the pool the allocator serves from is not the device default pool, that a freed block is still reserved after a free and a synchronize, that releasing returns it, that a release leaves a live allocation both reserved and usable, and that a negative index resolves to the current device. All four pass against the change and all four fail with the pool creation forced to fail, which is the behaviour without it. The live allocation test uses blocks of 64 MiB, because at a few megabytes two allocations share one driver reservation and nothing can be released while either is live. 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 have no equivalent in the HIP compatibility header, so both the allocator's pool code and these tests are compiled out on ROCm and the change is a no-op there. Not measured: Windows, and whether a second copy of this translation unit in one process creates a second pool. A duplicate pool would waste memory rather than misbehave, since each copy trims what it owns. --- backends/cuda/runtime/cuda_allocator.cpp | 181 +++++++++++++++++- backends/cuda/runtime/cuda_allocator.h | 37 ++++ backends/cuda/runtime/cuda_backend.cpp | 14 ++ .../cuda/runtime/test/test_cuda_allocator.cpp | 131 +++++++++++++ extension/cuda/runtime_api.h | 1 + 5 files changed, 363 insertions(+), 1 deletion(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 4c7d6aec288..68f5d0c6650 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -12,6 +12,13 @@ #include #include +#if !defined(EXECUTORCH_USE_HIP) +#include +#include +#include +#include +#endif + namespace executorch::backends::cuda { using executorch::runtime::Error; @@ -21,6 +28,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,7 +405,20 @@ Result CudaAllocator::allocate_async( DeviceIndex index, cudaStream_t stream) { void* ptr = nullptr; - cudaError_t err = cudaMallocAsync(&ptr, nbytes, stream); + cudaError_t err = cudaErrorInvalidValue; +#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 (pool != nullptr) { + err = cudaMallocFromPoolAsync(&ptr, nbytes, pool, stream); + } else { + err = cudaMallocAsync(&ptr, nbytes, stream); + } +#endif if (err != cudaSuccess) { ET_LOG( Error, @@ -313,6 +428,7 @@ Result CudaAllocator::allocate_async( static_cast(index)); return Error::MemoryAllocationFailed; } + return ptr; } @@ -323,6 +439,7 @@ void CudaAllocator::deallocate_async( if (ptr == nullptr) { return; } + cudaError_t err = cudaFreeAsync(ptr, stream); if (err != cudaSuccess) { ET_LOG( @@ -331,7 +448,69 @@ void CudaAllocator::deallocate_async( cudaGetErrorString(err), ptr, static_cast(index)); + return; + } +} + +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) { + continue; + } + // Only frees the driver has already observed can be released, so a caller + // that has not synchronized gets less back. That is why this does not + // synchronize itself: a device wide barrier here would wait on every stream + // on the device, including work this backend never queued. + const cudaError_t err = cudaMemPoolTrimTo(pool, 0); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolTrimTo failed for device %d: %s.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } } +#endif } Error CudaAllocator::memcpy_async( diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index b0a76a51f6d..4c802a7e00b 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -68,6 +68,43 @@ 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 this never affects memory another user of the async allocator is + * holding. + * + * Does nothing on ROCm, where the pool calls have no equivalent. + * + * @param index Device to release on, or a negative value to release every + * device this backend has allocated on. 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. + * + * @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 29cba8b5ada..39499638f78 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..02ba9c909a0 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -181,3 +181,134 @@ 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 means whichever device is current, which the allocator +// resolves rather than passing on to the driver. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsCurrentDeviceSentinel) { + 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); +} + +#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;