zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-06 03:58:12 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix + comment

Browse Source
This commit is contained in:
Awni Hannun
2025-11-04 11:29:37 -08:00
parent 7eaa504c26
commit e32db91988
2 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -106,6 +106,10 @@ CudaAllocator::CudaAllocator()
}
Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
if (size == 0) {
return Buffer{new CudaBuffer{nullptr, 0, -1}};
}
// Find available buffer from cache.
auto orig_size = size;
std::unique_lock lock(mutex_);
@@ -173,6 +177,10 @@ void CudaAllocator::free(Buffer buffer) {
if (!buf) {
return;
}
if (buf->size == 0) {
delete buf;
return;
}
std::unique_lock lock(mutex_);
active_memory_ -= buf->size;

4
mlx/backend/cuda/fence.cpp
View File

@@ -34,8 +34,8 @@ void Fence::update(Stream s, const array& a, bool cross_device) {
cbuf.device = -1;
auto& encoder = cu::device(s.device).get_command_encoder(s);
encoder.commit();
CHECK_CUDA_ERROR(
cudaMemcpyAsync(new_data, cbuf.data, cbuf.size, cudaMemcpyDefault));
CHECK_CUDA_ERROR(cudaMemcpyAsync(
new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream()));
CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream()));
cbuf.data = new_data;
}