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

fix memory count bug (#2717)

Browse Source
This commit is contained in:
Awni Hannun
2025-10-30 14:27:15 -07:00
committed by GitHub
parent 793a31eeb6
commit 68c5fa1c95
2 changed files with 15 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -97,7 +97,6 @@ CudaAllocator::CudaAllocator()
Buffer CudaAllocator::malloc(size_t size) {
// Find available buffer from cache.
auto orig_size = size;
std::unique_lock lock(mutex_);
if (size <= small_block_size) {
size = 8;
@@ -131,7 +130,7 @@ Buffer CudaAllocator::malloc(size_t size) {
}
lock.lock();
}
active_memory_ += size;
active_memory_ += buf->size;
peak_memory_ = std::max(active_memory_, peak_memory_);
// Maintain the cache below the requested limit.