[BUG][CUDA] Completion workers consume CPU while idle after the first batch #4451

☑️ I understand it is strictly prohibited to use AI to write issues.

Describe the bug
The CUDA completion worker may fail to go back to sleep after finishing the first batch and instead spin in a busy-loop, causing each worker tied to a relevant CUDA stream to pin a CPU core for an extended period.

The root cause is that Worker::thread_fn() in mlx/backend/cuda/worker.cpp initializes current_batch inside the loop body. It is reset to 0 on every iteration; once signaled_batch_ becomes greater than 0, the condition-variable wait predicate stays permanently true, so cond_.wait() can never re-block.

This leads to high idle CPU usage but does not affect CUDA computation correctness.

To Reproduce

In a CUDA environment, create multiple GPU streams and perform a synchronization on each stream to wake up its completion worker. Then let the process idle and measure the process CPU time during that window.

import time
import mlx.core as mx

streams = [mx.new_stream(mx.gpu) for _ in range(8)]

for stream in streams:
    mx.synchronize(stream)

cpu_before = time.process_time()
time.sleep(0.25)
cpu_seconds = time.process_time() - cpu_before

print(f"CPU time during idle wait: {cpu_seconds:.6f} seconds")

On a version containing the bug, the 8 workers idle for 250 ms and yield roughly 2.109 s of process CPU time, indicating that multiple workers are busy-looping in parallel.

Running the same test with the fixed version yields 0 s of process CPU time.

Expected behavior
After processing the current batch, if no new task arrives and no stop request has been received, the CUDA completion worker should block on the condition variable.

current_batch should persist across loop iterations to record the already-consumed batch number, so that the wait predicate becomes false again once the current batch is finished. An idle worker must not continuously scan an empty task queue or consume a CPU core.

Desktop (please complete the following information):

• OS：Microsoft Windows 10.0.26200.9168
• Clean MLX commit：43d2f06cb87e76895bf9a152bade4fee83408643
• Candidate commit：7b4e49587215c747bf2d2f553639f77af956fe33
• GPU：NVIDIA GeForce RTX 4060 Ti
• Compute capability：8.9
• NVIDIA driver：591.86
• CUDA toolkit：13.0.48
• Compiler：MSVC 19.44.35207，x64
• CMake：4.4.0
• Generator：Visual Studio 17 2022
• Build type：Release
• CUDA architecture：89

Additional context
Proposed fix:

• Move current_batch outside the while loop so it records the already-consumed batch number.
• Read stop_ while holding mtx_ to eliminate the data race caused by reading outside the lock and writing inside the lock.
• Keep task processing, exit, and drain behavior unchanged.

Regression test uses 8 CUDA streams and checks process CPU time during a 250 ms idle window:

• Before fix: ~2.109 s, test fails.
• After fix: 0 s, test passes.
• Threshold: < 0.1 s.
• Targeted test results: 1/1 test case, 2/2 assertions passed.

The fix branch has been rebased onto the latest upstream/main at the time; the corresponding fix commit is 7b4e49587.