zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

patch small pool init

Browse Source
This commit is contained in:
Awni Hannun
2025-07-29 13:57:39 -07:00
parent ef631d63af
commit 4f6bcf086e
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -30,8 +30,11 @@ SmallSizePool::SmallSizePool() {
next_free_ = buffer_; next_free_ = buffer_;
CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size)); CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));
CHECK_CUDA_ERROR( auto status =
cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetReadMostly, 0)); cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetReadMostly, 0);
if (status != cudaSuccess && status != cudaErrorInvalidValue) {
throw std::runtime_error("Unable to initialize small allocator pool");
}
auto curr = next_free_; auto curr = next_free_;
for (size_t i = 1; i < num_blocks; ++i) { for (size_t i = 1; i < num_blocks; ++i) {