zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-31 16:21:27 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

add pool threshold

Browse Source
This commit is contained in:
Awni Hannun
2025-10-27 16:26:57 -07:00
parent 58ccbaaf12
commit 21453281a4
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/backend/cuda/allocator.cpp
View File

@@ -99,8 +99,10 @@ CudaAllocator::CudaAllocator()
loc.id = 0; loc.id = 0;
loc.type = cudaMemLocationTypeNone; loc.type = cudaMemLocationTypeNone;
cudaMemGetDefaultMemPool(&cuda_pool_, &loc, cudaMemAllocationTypeManaged); cudaMemGetDefaultMemPool(&cuda_pool_, &loc, cudaMemAllocationTypeManaged);
// TODO set that. // TODO need a strategy for that
// uint64_t threshold = UINT64_MAX; uint64_t threshold = UINT64_MAX;
cudaMemPoolSetAttribute(
cuda_pool_, cudaMemPoolAttrReleaseThreshold, &threshold);
#endif #endif
} }