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

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

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

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