zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-06 12:09:43 +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) { 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. // Find available buffer from cache.
auto orig_size = size; auto orig_size = size;
std::unique_lock lock(mutex_); std::unique_lock lock(mutex_);
@@ -173,6 +177,10 @@ void CudaAllocator::free(Buffer buffer) {
if (!buf) { if (!buf) {
return; return;
} }
if (buf->size == 0) {
delete buf;
return;
}
std::unique_lock lock(mutex_); std::unique_lock lock(mutex_);
active_memory_ -= buf->size; 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; cbuf.device = -1;
auto& encoder = cu::device(s.device).get_command_encoder(s); auto& encoder = cu::device(s.device).get_command_encoder(s);
encoder.commit(); encoder.commit();
CHECK_CUDA_ERROR( CHECK_CUDA_ERROR(cudaMemcpyAsync(
cudaMemcpyAsync(new_data, cbuf.data, cbuf.size, cudaMemcpyDefault)); new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream()));
CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream())); CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream()));
cbuf.data = new_data; cbuf.data = new_data;
} }