fix(rocm): initialize Device's rocBLAS handle once, under a lock #2198

Part of #1801

Problem / Background

The ROCm overlay's Device creates its rocBLAS handle lazily and without a lock. Two threads that run their first matmul at the same time can get back a null handle, and every later GEMM rebinds one shared handle to the calling thread's stream with no lock either, so a GEMM can be enqueued on another thread's stream. The review of #2197 (lock the rocm::device() map) found the init race and kept it out of that issue. Neither failure has been observed yet.

Current Behavior

All paths are under src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/, read on main at 84d3a7b. The Device class has these lazy inits, all with the same flag-before-work pattern and no lock:

• get_rocblas_handle() (device.cpp:175-245) sets rocblas_initialized_ = true at :177, then does the arch check and rocblas_create_handle(&rocblas_) at :227. A second thread that arrives in that window sees the flag, sees rocblas_available_ still at its default true (device.h:498), and returns rocblas_, which is still nullptr. If both threads read the flag before either writes it, both create a handle and one leaks. On a create failure, or an unsupported arch, it prints a warning and keeps rocblas_available_ = false forever.
• set_rocblas_stream() (device.cpp:431-436) compares and writes rocblas_stream_ and calls rocblas_set_stream on the one shared handle, unlocked. Every rocBLAS call site does set_rocblas_stream(stream) and then issues its GEMM, so thread B's rebind between thread A's rebind and A's GEMM puts A's GEMM on B's stream, where A's stream does not wait for it.
• is_rocblas_bf16_available() (device.cpp:257-351) sets rocblas_bf16_probed_ first, runs a probe GEMM, and on failure calls hipDeviceReset() and destroys and recreates rocblas_. Nothing calls it (no caller in the tree outside device.cpp/device.h).
• has_native_wmma() (:353-397) and supports_cdna_mfma_gemm() (:399-418) set wmma_probed_ / cdna_mfma_probed_ before computing the result, so a concurrent caller reads false and takes the non-WMMA route once. has_native_wmma() is called from scaled_dot_product_attention.cpp:181,237,307 and quantized/qmm.hip:898,6087. supports_cdna_mfma_gemm() has no caller.

The flags are plain bools read and written from several threads, which is a data race. The rocBLAS call sites are the launch_kernel lambdas at matmul.cpp:241,499,873, gemms/rocblas_gemm.cpp:80,299,436 and quantized/qmm.hip:1278,1521, plus matmul.cpp:223,486, which fetch the handle outside the lambda. CommandEncoder::launch_kernel (device.h:574) runs the lambda synchronously on the calling thread. Several sites discard the GEMM status (for example the f32 rocblas_sgemm at matmul.cpp:297), so a null handle leaves the output unwritten without an error.

Upstream MLX at the pin (81ba1c6a, mlx/backend/cuda/cublas_utils.cpp:44-79) keeps cuBLASLt handles in static thread_local std::vector<CublasHandles> cache(gpu::device_count()), so each thread creates its own handle per device and nothing is shared. Creation errors throw through CHECK_CUBLAS_ERROR. The stream is passed per call to cublasLtMatmul (cublas_utils.cpp:199-215), not bound on the handle.

Reachability

Real on one GPU today, though narrow and not observed. is_rocblas_available() runs on the first matmul that is not handled by hipBLASLt (matmul.cpp:775, every f32 GEMM) and on every affine quantized matmul route decision (quantized/qmm.hip:876), so the first quantized model forward in a process creates the handle. The server runs forwards on several OS threads on the same Device, each with its own thread-local stream and no global eval lock: BatchScheduler, EmbeddingWorker (src/server/embedding_worker.rs:132), RerankWorker (src/server/rerank_worker.rs:113) and AudioWorker (src/server/audio_worker.rs:165), as established in #2183. With -m <chat> --embedding-model <emb>, first requests to both models that arrive together race the init; there is no startup warmup forward that would create the handle first. The stream rebinding race is not limited to first use: it is open on every rocBLAS GEMM while two of those threads run rocBLAS GEMMs. Priority is medium: reachable in a supported single-GPU server configuration, and the failure is a silently wrong GEMM output, but the init window is once per process and nobody has hit it.

Proposed Solution

One rocBLAS handle per Device, created once under a mutex, and every rocBLAS call made while holding that mutex through a lease that binds the stream.

• device.h: add std::mutex rocblas_mtx_ and std::atomic<bool> rocblas_ready_{false}. Replace rocblas_initialized_ and rocblas_available_ with bool rocblas_arch_supported_ (computed in the constructor) and bool rocblas_warned_{false}. Keep rocblas_ and rocblas_stream_, now guarded by rocblas_mtx_.
• Add a move-only class RocblasLease in device.h holding a std::unique_lock<std::mutex> and the rocblas_handle, with rocblas_handle handle() const. Add RocblasLease Device::acquire_rocblas(hipStream_t stream): lock rocblas_mtx_, create the handle if needed (below), call rocblas_set_stream(handle, stream) when rocblas_stream_ != stream and throw std::runtime_error("rocblas_set_stream failed (status N)") on failure, then update rocblas_stream_ and return the lease.
• Handle creation, done with rocblas_mtx_ held: if !rocblas_arch_supported_, throw the existing "rocBLAS is not available on this GPU architecture" error. Otherwise make_current(), create into a local rocblas_handle, and on failure throw std::runtime_error naming the status, leaving rocblas_ready_ false so the next call retries. On success store rocblas_ and only then rocblas_ready_.store(true, std::memory_order_release). The first failure of either kind also prints today's stderr warning once (rocblas_warned_).
• is_rocblas_available(): return false if !rocblas_arch_supported_; return true if rocblas_ready_ (acquire load) without locking; otherwise lock, try to create, return whether it succeeded. Because a ready device never locks here, it is safe to call while holding a lease.
• Make get_rocblas_handle() private (it becomes the locked create helper) and delete set_rocblas_stream(), so an unguarded rocBLAS call no longer compiles. Convert the eight lambdas listed above: take auto lease = device.acquire_rocblas(stream); immediately before the first rocBLAS call and use lease.handle() for every call in the lambda. In quantized/qmm.hip:1278 acquire it after the hipBLASLt attempt, so hipBLASLt GEMMs do not hold the rocBLAS lock. Delete the outer handle fetches at matmul.cpp:223,486.
• Compute rocblas_arch_supported_, has_native_wmma_ and cdna_mfma_ok_ once in Device::Device from the hipDeviceProp_t it already reads (device.cpp:118-130), with the existing arch lists, and make has_native_wmma() and supports_cdna_mfma_gemm() plain getters. Delete wmma_probed_ and cdna_mfma_probed_.
• Delete is_rocblas_bf16_available() and its members rocblas_bf16_probed_ / rocblas_bf16_available_. It has no caller, and its recovery path (hipDeviceReset() plus recreating the shared handle) cannot be made safe while other threads use the device.
• Rejected: per-thread handles as upstream does. rocBLAS binds the stream on the handle rather than per call, the fork is built around one handle per device, and per-thread handles add a handle and its device workspace per worker thread and move creation onto each new thread's first GEMM, possibly inside a decode-step stream capture where rocblas_create_handle has not been checked. The lease serializes only host-side enqueue; GEMMs on different streams still run concurrently on the GPU.

Scope

In scope: device.h, device.cpp, the rocBLAS call sites in matmul.cpp, gemms/rocblas_gemm.cpp and quantized/qmm.hip, a new test file, and a LOCAL_FIXES entry.

Out of scope, to be filed separately: gemms/hipblaslt_gemm.cpp has the same init pattern (initialized read outside the mutex at :89 and :1044 and set before hipblasLtCreate at :58, so a racing caller gets "hipBLASLt is not available" and falls back to rocBLAS), an unlocked ensure_workspace (:99-126), one per-device workspace shared by concurrent hipblasLtMatmul calls on different streams, and a GemmPipe* used outside g_pipe_mutex (:404-449) that another thread may erase. Also out of scope: the CUDA and Metal backends.

Implementation Notes

• Reuse: model the test on tests/rocm_qmm_env.rs: a parent test that spawns std::env::current_exe() with <child> --exact --ignored, a CHILD_ENV marker, extra env vars, and a 120 s deadline (run_child, :213-283), plus an #[ignore] child that returns early without the marker. Use per-thread streams and a std::sync::Barrier as in tests/rocm_jit_module_concurrency.rs (PR fix(rocm): lock the JIT module cache and the HIP event pool #2196): streams::new_thread_local_generation_stream and install_thread_local_default_stream.
• Edge cases: two threads creating at once (one creation, both get the same handle); create failure (throws, not cached, next call retries); unsupported arch (permanent, no lock after the first check); a lease held while is_rocblas_available() runs (no deadlock); an exception inside a lambda (the lease unlocks on unwind).
• Error handling: callers that check is_rocblas_available() keep their naive-GEMM fallback; a direct acquire_rocblas on a device without rocBLAS throws as today.
• LOCAL_FIXES: add an entry to src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md under "Fixes to the fork's kernels" with the next free number, naming mlx/backend/rocm/device.cpp, device.h, matmul.cpp, gemms/rocblas_gemm.cpp and quantized/qmm.hip, stating what was unlocked and how it is now locked, and ending with "Applies to the fork; kept in mlxcelverse under the 2026-10-06 fork policy, not proposed there." Add nothing under docs/mlxcelverse/upstream/.

Acceptance Criteria

• Device creates the rocBLAS handle under rocblas_mtx_, sets rocblas_ready_ only after a successful create, and throws with the rocBLAS status on failure; rocblas_initialized_, rocblas_available_, set_rocblas_stream, is_rocblas_bf16_available, wmma_probed_ and cdna_mfma_probed_ are gone.
• Every rocBLAS call in the overlay goes through acquire_rocblas(stream) (grep -rn "get_rocblas_handle\|rocblas_set_stream" src/lib/mlx-cpp/patches-rocm matches only device.cpp and device.h).
• tests/rocm_rocblas_handle_concurrency.rs (#![cfg(feature = "rocm")], skips when gpu_backend_kind() != GpuBackendKind::Rocm) has a parent test first_gemm_from_many_threads_in_fresh_processes that runs the child 10 times, each in a fresh process with MLX_NO_HIPBLASLT=1 so f32 GEMMs take rocBLAS. The child child_first_gemm_from_many_threads starts 8 threads; each installs its own stream, waits on the barrier with no GPU matmul run yet in the process, then computes 33 f32 matmuls of [64, 128] x [128, 64] with integer inputs in -4..=4 from a per-thread seed, and asserts each output exactly equals the host i32 reference.
• With the lock removed (the old device.cpp restored), run the test file 20 times and report in the PR how many runs failed. If none did, say so: the test is then a guard, not a deterministic reproduction.
• The LOCAL_FIXES entry exists with no upstream marking, and make verify-rocm-overlay passes.
• Existing ROCm GEMM, SDPA and qmm tests pass under make verify-rocm.

Verification

cargo test --profile test-fast --features rocm --test rocm_rocblas_handle_concurrency
cargo clippy --features rocm --test rocm_rocblas_handle_concurrency -- -D warnings
grep -rn "get_rocblas_handle\|rocblas_set_stream\|set_rocblas_stream" src/lib/mlx-cpp/patches-rocm
make verify-rocm-overlay
make verify-rocm

A pass is all ten child runs exact, the grep limited to device.cpp/device.h, and both make targets green on an AMD host.