From 7b4e49587215c747bf2d2f553639f77af956fe33 Mon Sep 17 00:00:00 2001 From: YH Yan Date: Tue, 1 Sep 2026 12:09:05 +0800 Subject: [PATCH 1/2] [CUDA] Fix completion worker busy loop --- mlx/backend/cuda/worker.cpp | 10 +++++--- tests/CMakeLists.txt | 5 ++++ tests/gpu_tests.cpp | 50 +++++++++++++++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 3 deletions(-) diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp index de38a06e51..4333618abf 100644 --- a/mlx/backend/cuda/worker.cpp +++ b/mlx/backend/cuda/worker.cpp @@ -57,14 +57,18 @@ void Worker::commit(cudaStream_t stream) { } void Worker::thread_fn() { - while (!stop_) { - uint64_t current_batch = 0; + // Keep the consumed batch count across iterations so the wait predicate + // becomes false again after a completed batch. + uint64_t current_batch = 0; + bool stop = false; + while (!stop) { Tasks tasks; { std::unique_lock lk(mtx_); - cond_.wait(lk, [this, ¤t_batch] { + cond_.wait(lk, [this, current_batch] { return this->signaled_batch_ > current_batch || this->stop_; }); + stop = stop_; current_batch = signaled_batch_; auto end = worker_tasks_.upper_bound(current_batch); for (auto it = worker_tasks_.begin(); it != end; ++it) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 859215acb1..75a152f772 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,6 +41,11 @@ target_sources( ${METAL_TEST_SOURCES}) target_link_libraries(tests PRIVATE mlx doctest) + +if(MLX_BUILD_CUDA) + target_compile_definitions(tests PRIVATE MLX_TEST_CUDA) +endif() + target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS}) target_link_options(tests PRIVATE ${SANITIZER_LINK_FLAGS}) diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index 8f492e52f4..9fad3756e5 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -3,16 +3,66 @@ #include #include #include +#include #include +#include #include "doctest/doctest.h" #include "mlx/mlx.h" +#if defined(MLX_TEST_CUDA) && defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#endif + using namespace mlx::core; static const std::array types = {bool_, uint32, int32, int64, float32}; +#ifdef MLX_TEST_CUDA +namespace { +double process_cpu_seconds() { +#ifdef _WIN32 + FILETIME creation; + FILETIME exit; + FILETIME kernel; + FILETIME user; + if (!GetProcessTimes(GetCurrentProcess(), &creation, &exit, &kernel, &user)) { + return -1.0; + } + ULARGE_INTEGER kernel_ticks; + kernel_ticks.LowPart = kernel.dwLowDateTime; + kernel_ticks.HighPart = kernel.dwHighDateTime; + ULARGE_INTEGER user_ticks; + user_ticks.LowPart = user.dwLowDateTime; + user_ticks.HighPart = user.dwHighDateTime; + return static_cast(kernel_ticks.QuadPart + user_ticks.QuadPart) / 1e7; +#else + return static_cast(std::clock()) / CLOCKS_PER_SEC; +#endif +} +} // namespace + +TEST_CASE("test cuda workers do not consume cpu while idle") { + std::vector streams; + for (int i = 0; i < 8; ++i) { + streams.push_back(new_stream(Device::gpu)); + synchronize(streams.back()); + } + + auto before = process_cpu_seconds(); + REQUIRE(before >= 0.0); + std::this_thread::sleep_for(std::chrono::milliseconds(250)); + auto cpu_seconds = process_cpu_seconds() - before; + + INFO("process CPU seconds during 250 ms idle wait: " << cpu_seconds); + CHECK(cpu_seconds < 0.1); +} +#endif + TEST_CASE("test gpu arange") { for (auto t : types) { if (t == bool_) { From b7a047e47b84bb3bad6245bf3740d676885df4b1 Mon Sep 17 00:00:00 2001 From: YH Yan Date: Sat, 5 Sep 2026 17:01:53 +0800 Subject: [PATCH 2/2] Simplify worker shutdown control flow. Remove potentially flaky CPU timing tests. --- mlx/backend/cuda/worker.cpp | 9 +++---- tests/CMakeLists.txt | 5 ---- tests/gpu_tests.cpp | 50 ------------------------------------- 3 files changed, 4 insertions(+), 60 deletions(-) diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp index 4333618abf..6266d6099b 100644 --- a/mlx/backend/cuda/worker.cpp +++ b/mlx/backend/cuda/worker.cpp @@ -57,18 +57,17 @@ void Worker::commit(cudaStream_t stream) { } void Worker::thread_fn() { - // Keep the consumed batch count across iterations so the wait predicate - // becomes false again after a completed batch. uint64_t current_batch = 0; - bool stop = false; - while (!stop) { + while (true) { Tasks tasks; { std::unique_lock lk(mtx_); cond_.wait(lk, [this, current_batch] { return this->signaled_batch_ > current_batch || this->stop_; }); - stop = stop_; + if (stop_) { + return; + } current_batch = signaled_batch_; auto end = worker_tasks_.upper_bound(current_batch); for (auto it = worker_tasks_.begin(); it != end; ++it) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 75a152f772..859215acb1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,11 +41,6 @@ target_sources( ${METAL_TEST_SOURCES}) target_link_libraries(tests PRIVATE mlx doctest) - -if(MLX_BUILD_CUDA) - target_compile_definitions(tests PRIVATE MLX_TEST_CUDA) -endif() - target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS}) target_link_options(tests PRIVATE ${SANITIZER_LINK_FLAGS}) diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index 9fad3756e5..8f492e52f4 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -3,66 +3,16 @@ #include #include #include -#include #include -#include #include "doctest/doctest.h" #include "mlx/mlx.h" -#if defined(MLX_TEST_CUDA) && defined(_WIN32) -#ifndef NOMINMAX -#define NOMINMAX -#endif -#include -#endif - using namespace mlx::core; static const std::array types = {bool_, uint32, int32, int64, float32}; -#ifdef MLX_TEST_CUDA -namespace { -double process_cpu_seconds() { -#ifdef _WIN32 - FILETIME creation; - FILETIME exit; - FILETIME kernel; - FILETIME user; - if (!GetProcessTimes(GetCurrentProcess(), &creation, &exit, &kernel, &user)) { - return -1.0; - } - ULARGE_INTEGER kernel_ticks; - kernel_ticks.LowPart = kernel.dwLowDateTime; - kernel_ticks.HighPart = kernel.dwHighDateTime; - ULARGE_INTEGER user_ticks; - user_ticks.LowPart = user.dwLowDateTime; - user_ticks.HighPart = user.dwHighDateTime; - return static_cast(kernel_ticks.QuadPart + user_ticks.QuadPart) / 1e7; -#else - return static_cast(std::clock()) / CLOCKS_PER_SEC; -#endif -} -} // namespace - -TEST_CASE("test cuda workers do not consume cpu while idle") { - std::vector streams; - for (int i = 0; i < 8; ++i) { - streams.push_back(new_stream(Device::gpu)); - synchronize(streams.back()); - } - - auto before = process_cpu_seconds(); - REQUIRE(before >= 0.0); - std::this_thread::sleep_for(std::chrono::milliseconds(250)); - auto cpu_seconds = process_cpu_seconds() - before; - - INFO("process CPU seconds during 250 ms idle wait: " << cpu_seconds); - CHECK(cpu_seconds < 0.1); -} -#endif - TEST_CASE("test gpu arange") { for (auto t : types) { if (t == bool_) {