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

[CUDA] Check CUDA error in synchronize (#2757)

Browse Source
This commit is contained in:
Cheng
2025-11-14 07:10:23 +09:00
committed by GitHub
parent 66519fb348
commit b704e9e77a
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/cuda/device.cpp
View File

@@ -368,7 +368,7 @@ void CommandEncoder::commit() {
} }
void CommandEncoder::synchronize() { void CommandEncoder::synchronize() {
cudaStreamSynchronize(stream_); CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
auto p = std::make_shared<std::promise<void>>(); auto p = std::make_shared<std::promise<void>>();
std::future<void> f = p->get_future(); std::future<void> f = p->get_future();
add_completed_handler([p = std::move(p)]() { p->set_value(); }); add_completed_handler([p = std::move(p)]() { p->set_value(); });

2
mlx/backend/cuda/worker.cpp
View File

@@ -44,7 +44,7 @@ void Worker::commit(cudaStream_t stream) {
} }
signal_event_.record(stream); signal_event_.record(stream);
signal_event_.wait(signal_stream_); signal_event_.wait(signal_stream_);
cudaLaunchHostFunc(signal_stream_, signal, this); CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this));
} }
void Worker::thread_fn() { void Worker::thread_fn() {