fix(rocm): make hipBLASLt GEMM state safe across threads and streams #2200

Part of #1801

Problem / Background

The review of #2198 (rocBLAS handle init and stream rebind races in the ROCm overlay Device) found the same class of race in the hipBLASLt GEMM layer, which lives outside Device and is out of that issue's scope. hipBLASLt is the default route for every bf16/fp16 GEMM, because use_hip_graphs() always returns false (device.cpp:43-54) and the graph gate in is_hipblaslt_available() (gemms/hipblaslt_gemm.cpp:1036) never trips. The server runs forwards on several OS threads that share one Device, each on its own stream: BatchScheduler, EmbeddingWorker, RerankWorker and AudioWorker (#2183, #2196). So -m <chat> --embedding-model <emb> puts two threads into this file at once, and the opt-in async MoE path (MLX_ROCM_MOE_ASYNC=1, matmul.cpp:53-100, :1454-1462) adds a third caller: the HIP host-callback thread, which serves every stream.

Current Behavior

Paths are under src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/hipblaslt_gemm.cpp, read on main at 84d3a7b. Shared mutable state:

Failure modes, worst first:

• Shared workspace (corrupts output). Every hipblasLtMatmul on the device gets the same 32 MB buffer (:417, :858, :986, :1633, and the tune loops at :777 and :1562). The call only enqueues, so two GEMMs on two streams whose algorithms request workspace run on the GPU at the same time over the same scratch memory, and both outputs are wrong. Stream order protects reuse only within one stream.
• Init window, worse interleaving (corrupts output or faults). Thread A sets initialized (:58), creates the handle, sets available (:68), and only then allocates the workspace (:78-83). Thread B arriving after :68 skips the lock, runs a GEMM needing workspace, sees workspace == nullptr (:106), and calls hipMalloc(&state.workspace, required) (:118) unlocked while A writes the same fields. The ordering A ptr, B ptr, B size (:124), A size (:79) leaves {B's small buffer, 32 MB}, and later GEMMs with a larger workspace write past B's allocation. Other orderings leak one buffer.
• Workspace growth (GPU use-after-free). When the recorded size is below a request (init allocation failed at :78, or the interleaving above), ensure_workspace frees the buffer (:111) that another stream's queued GEMM still uses.
• Pipe erase (host use-after-free). On a failed cached-pipe matmul (stale bucketed algo, the case :441-442 describes), the thread destroys the pipe's layouts and desc (:447-451) and erases it while another thread holds pipe from :409, or use_la..use_desc from :905-923, and passes them to hipblasLtMatmul. The erase is reachable: the retry at :949-1008 fixes algo_cache but leaves the bad algo in the pipe, so the next call on that key fails and erases. The comment at :295 ("Process-lifetime; never destroyed") is stale.
• Init window, common interleaving (wrong route, no corruption). Thread B arriving between :58 and :68 reads initialized == true, skips the lock, and sees available == false: is_hipblaslt_available() returns false and that GEMM falls back to rocBLAS, and get_handle() throws "hipBLASLt is not available on this device." (callers such as matmul.cpp:199 catch and fall back). If B is inside gemm_caps, the throw leaves g_caps[d] all false forever, so device_has_fp8_gemm() (:1426) disables the fp8 prefill route (quantized/qmm.hip:884) for the process.
• init_handle calls hipblasLtCreate without making device_id current (hipSetDevice wraps only the hipMalloc, :75-84).

Measured on the only AMD host available (gfx1151, ROCm 7.15, hipBLASLt 1.4): with HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES at 32 MB, hipblasLtMatmulAlgoGetHeuristic returned workspaceSize == 0 for all 8 candidates of each of 15 shapes (bf16 TN decode, prefill and MLP shapes from [4096, 1] x 4096 to [11008, 512] x 4096 and K up to 1M; f32 NN up to 2048^3). So on that device the three workspace races never run, and the reachable races are the pipe erase and the init window. Archs whose Tensile libraries pick split-K algorithms with global scratch (CDNA, gfx1201) were not measured.

Upstream MLX at the pin (81ba1c6a): mlx/backend/cuda/cublas_utils.cpp:44-79 keeps a cuBLASLt handle and preference per thread per device in static thread_local std::vector<CublasHandles>, creation throws through CHECK_CUBLAS_ERROR, and the stream is passed per call (:197-215). The workspace is allocated per call by allocate_workspace (mlx/backend/cuda/utils.cpp:82-99): cu::malloc_async from the MLX allocator, rounded up to 256 bytes, registered with encoder.add_temporary so it lives until that command buffer completes, with a debug assert that the stream is not capturing.

Proposed Solution

All changes are in gemms/hipblaslt_gemm.cpp and gemms/hipblaslt_gemm.h, plus one call in device.cpp.

Handle init. Replace initialized and available with std::atomic<int> init_state{0} (0 not tried, 1 ready, 2 unavailable). Add void ensure_handle(HipblasltState&, int device_id): return on a nonzero acquire load; otherwise lock state.mutex, recheck, save and set the current device, hipblasLtCreate into a local, restore the device, then on success store state.handle and init_state.store(1, std::memory_order_release); on failure print today's warning and store 2. Failure stays permanent as today, because is_hipblaslt_available() runs on every GEMM route decision and must not retry a create each time. get_handle() and is_hipblaslt_available() call ensure_handle and then read init_state; neither reads a field outside the lock before the flag says ready. Remove the workspace allocation from init. Lock order stays g_caps_mutex then state.mutex. Keep one handle per device. Rejected: per-thread handles as upstream does, because algo_cache and g_pipe_cache hold heuristic results computed on one handle and shared across threads, and the async MoE path would create handles inside a HIP host callback.

Workspace: one per stream. Replace ensure_workspace and the two HipblasltState workspace fields with std::pair<void*, size_t> stream_workspace(int device_id, hipStream_t stream, size_t required) over static std::mutex g_ws_mutex and static std::unordered_map<hipStream_t, void*> g_ws[kMaxDevices]. required == 0 returns {nullptr, 0} without locking. Otherwise, under the lock, return the stream's buffer with size kMaxWorkspaceBytes; if it has none, allocate exactly kMaxWorkspaceBytes with device_id current and store it. Every heuristic preference already caps at that size (:186, :545-550, :1485), so a buffer never grows and never moves while its stream lives. Under stream capture (hipStreamGetCaptureInfo not None) with no buffer yet, throw std::runtime_error("hipBLASLt: algorithm needs N bytes of workspace on a capturing stream with none allocated") rather than call hipMalloc; today's capture-capable hipBLASLt call is the fp8 path (:1550-1553), which requires M >= 64 and is not reached from decode capture. On allocation failure return {nullptr, 0}: the tune loops keep skipping that algo (:780, :1565), and every other caller (:417, :858, :986, :1633) throws the existing "hipBLASLt: failed to allocate workspace of N bytes" message instead of passing a null buffer to hipblasLtMatmul. Add void hipblaslt_release_stream_workspace(hipStream_t stream) to hipblaslt_gemm.h, which synchronizes the stream, frees its buffer and erases the entry, and call it from CommandEncoder::~CommandEncoder (device.cpp:477) so a later stream that reuses the handle value cannot inherit a buffer its predecessor is still draining. Cost: 32 MB per stream that runs a workspace-needing GEMM, and nothing on devices like gfx1151 whose heuristics request none (which also drops the 32 MB per device that init allocates today).

Why per-stream and not a lease: a lease held for the duration of the hipblasLtMatmul call ends when the call returns, before the kernel runs, so it does not stop the next lessee's kernel on another stream from overlapping on the GPU. Making it safe needs either hipStreamSynchronize inside the lease, which serializes GEMMs and is illegal under capture, or event-ordered handoff (hipStreamWaitEvent between unrelated workers' streams), which serializes GEMMs on the GPU. A thread-local buffer is also wrong, because the HIP callback thread runs the async MoE GEMMs for every stream. Per-call allocator temporaries as upstream does are not available here, because hipblaslt_gemm_raw, hipblaslt_gemm_fp8_raw and hipblaslt_gemm_rowmajor_on_stream receive a bare hipStream_t with no CommandEncoder, and a per-call hipMalloc is not capturable.

Pipe lifetime: shared ownership. Store std::shared_ptr<GemmPipe> in g_pipe_cache. Give GemmPipe a destructor that destroys its four layouts and desc, and delete its copy operations. The fast path copies the shared_ptr under g_pipe_mutex and uses it after unlock; the failure path at :443-454 only erases the map entry, and the descriptors die with the last holder. The miss path builds a std::make_shared<GemmPipe>, emplaces it, adopts the existing entry when !inserted, and holds that shared_ptr through the matmul and the retry (:929-1008). Delete the manual destroy calls at :447-451 and :914-918 and fix the comment at :292-295. Rejected: holding g_pipe_mutex across hipblasLtMatmul, which would serialize host enqueue of every hipBLASLt GEMM in the process, including the MoE segment loop.

Scope

In scope: gemms/hipblaslt_gemm.cpp, gemms/hipblaslt_gemm.h, the release call in device.cpp, a new test file, a LOCAL_FIXES entry, and removing MLX_NO_HIPBLASLT=1 from tests/rocm_rocblas_handle_concurrency.rs once #2198 has landed.

Out of scope: the rocBLAS handle and stream races (#2198), the leaked new float[2] alpha/beta pairs (:1093, :1161), the retry path not refreshing a pipe's stale algo (:949-1008, which only makes the erase reachable), and the CUDA and Metal backends.

Ordering: land #2198 first. Both touch adjacent GEMM code, every bf16 Matmul GEMM first passes is_rocblas_available() (matmul.cpp:775) before reaching hipBLASLt, and the #2198 test runs with MLX_NO_HIPBLASLT=1 because of this issue.

Implementation Notes

• Reuse: model the test on fix(rocm): initialize Device's rocBLAS handle once, under a lock #2198's tests/rocm_rocblas_handle_concurrency.rs and on tests/rocm_qmm_env.rs (run_child at :213-283: re-run std::env::current_exe() with <child> --exact --ignored, a CHILD_ENV marker, captured stderr, 120 s deadline). Per-thread streams and the barrier come from tests/rocm_jit_module_concurrency.rs (PR fix(rocm): lock the JIT module cache and the HIP event pool #2196): mlxcel_core::streams::new_thread_local_generation_stream and install_thread_local_default_stream. Build inputs with mlxcel_core::from_slice_f32, astype(.., dtype::BFLOAT16), matmul and transpose.
• Edge cases: two threads initializing at once (one create, both get the handle); create failure (warning once, state 2, no lock after); two streams needing workspace (distinct buffers); the same stream from the encoding thread and the HIP callback thread (same buffer, stream-ordered); a stream destroyed and its handle value reused (entry released first); a pipe erased while another thread uses it (descriptors live until that thread drops its shared_ptr).
• Error handling: callers that catch hipBLASLt exceptions keep their rocBLAS fallback (matmul.cpp:199, quantized/qmm.hip:1293); the workspace and capture errors above are std::runtime_errors naming the byte count.
• 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 at merge time, naming mlx/backend/rocm/gemms/hipblaslt_gemm.cpp, gemms/hipblaslt_gemm.h and device.cpp, stating what was shared or unlocked and how each is now owned, 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

• init_state is published with release after a successful hipblasLtCreate made with the target device current; no field of HipblasltState is read outside state.mutex before an acquire load sees it set.
• ensure_workspace is gone; every hipblasLtMatmul in the file takes its buffer from stream_workspace(device_id, stream, ..), two streams never share one, and ~CommandEncoder releases its stream's buffer.
• g_pipe_cache holds std::shared_ptr<GemmPipe>; no raw GemmPipe* or descriptor copied from a map entry is used after g_pipe_mutex is released without a shared_ptr keeping it alive. (Superseded by measurement in PR fix(rocm): make hipBLASLt GEMM state safe across threads and streams #2208: hipblasLtMatmul writes into the descriptor objects, so a pipe shared across threads fails however it is owned, shared_ptr included, 0 of 5 children; the pipe cache is thread-local instead, 5 of 5 and 20 of 20 runs.)
• tests/rocm_hipblaslt_concurrency.rs (#![cfg(feature = "rocm")], skips when gpu_backend_kind() != GpuBackendKind::Rocm) has a parent test hipblaslt_gemms_from_many_threads_in_fresh_processes that runs the #[ignore] child 10 times, each in a fresh process with MLX_NO_HIPBLASLT and MLX_HIPBLASLT_NO_PIPE_CACHE removed from its environment. The child child_hipblaslt_gemms_from_many_threads starts 8 threads; each installs its own stream and waits on a Barrier before any GPU matmul has run in the process, then runs 33 rounds of a bf16 [64, 64] x transpose([64, 64]) (hipBLASLt TN pipe), a bf16 batched [4, 64, 64] x [4, 64, 64] (hipblaslt_gemm_batched) and an f32 [64, 128] x [128, 64] (rocBLAS), all threads using the same shapes so they contend on the same pipe keys. Inputs are integers from a per-thread seed, in -2..=2 for bf16 (K = 64, so every output is an integer of magnitude at most 256 and exact in bf16) and -4..=4 for f32, and each output cast to f32 must equal the host i32 reference exactly. (bf16 inputs are -1..=1 in PR fix(rocm): make hipBLASLt GEMM state safe across threads and streams #2208: on gfx1151 both libraries return residues near 2^-22 at exact-zero outputs when bf16 inputs mix zeros with magnitude 2, fix(rocm): bf16 GEMMs return 2^-22 residue at exact-zero outputs on gfx1151 #2206.) The parent also asserts each child's stderr contains [hipBLASLt caps] device exactly once (printed by gemm_caps, :234-242), which proves the bf16 GEMMs reached hipBLASLt.
• With the fix removed (the old hipblaslt_gemm.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 reproduction. The PR also reports the workspaceSize values the heuristics return on the test host for the test shapes, so the reader knows whether the workspace path ran.
• tests/rocm_rocblas_handle_concurrency.rs no longer sets MLX_NO_HIPBLASLT and still passes.
• The LOCAL_FIXES entry exists with the fork-policy wording and no upstream marking, and make verify-rocm-overlay passes.
• Existing ROCm GEMM, MoE and qmm tests pass under make verify-rocm.

Verification

cargo test --profile test-fast --features rocm --test rocm_hipblaslt_concurrency
cargo test --profile test-fast --features rocm --test rocm_rocblas_handle_concurrency
cargo clippy --features rocm --test rocm_hipblaslt_concurrency -- -D warnings
grep -n "ensure_workspace\|GemmPipe\*\|initialized" src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/hipblaslt_gemm.cpp
make verify-rocm-overlay
make verify-rocm

A pass is all ten child runs exact with one caps line each, the grep empty, and both make targets green on an AMD host.

Priority

Medium. Output corruption needs an algorithm that requests workspace, and on gfx1151 every candidate measured requested none, so the shared-workspace and growth races do not run there. What is reachable on one GPU today is the pipe-erase use-after-free (needs a stale cached algo plus a concurrent same-geometry GEMM) and the init window, which only misroutes one GEMM or permanently disables the fp8 route. Raise to high if the PR's workspace report, or a run on gfx1201 or CDNA, shows nonzero workspaceSize for a shape the server runs.