From 21453281a4692409b50ff9e0c3b84bf01bf2d12b Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Mon, 27 Oct 2025 16:26:57 -0700
Subject: [PATCH] add pool threshold

---
 mlx/backend/cuda/allocator.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp
index 97a5ae4d4..6ef2da907 100644
--- a/mlx/backend/cuda/allocator.cpp
+++ b/mlx/backend/cuda/allocator.cpp
@@ -99,8 +99,10 @@ CudaAllocator::CudaAllocator()
   loc.id = 0;
   loc.type = cudaMemLocationTypeNone;
   cudaMemGetDefaultMemPool(&cuda_pool_, &loc, cudaMemAllocationTypeManaged);
-  // TODO set that.
-  // uint64_t threshold = UINT64_MAX;
+  // TODO need a strategy for that
+  uint64_t threshold = UINT64_MAX;
+  cudaMemPoolSetAttribute(
+      cuda_pool_, cudaMemPoolAttrReleaseThreshold, &threshold);
 #endif
 }