fix(rocm): set blocking-sync before first HIP queue, FFT cache to 128

docs/environment-variables.md

@@ -189,10 +189,11 @@ CUDA builds also use non-`MLXCEL_*` variables such as `CUDA_HOME` and
 ROCm builds (`--features rocm`) read `ROCM_PATH` (default `/opt/rocm`) and
 `MLX_ROCM_ARCHITECTURES` (the `gfx` targets, default detected with `rocminfo`); see
 [Installation](installation.md#linux-with-amd-rocm-experimental). At runtime they
-also read `MLX_ROCM_FFT_CACHE_SIZE` (default 8), the number of hipFFT plans kept
-alive. The default is low on purpose: past roughly a dozen live plans, creating
-the next one blocks inside hipFFT with the GPU idle (lablup/mlxcel#1825). Raise
-it only while measuring that limit. `MLX_ROCM_GPU_WATCHDOG_SECS` (default `0`,
+also read `MLX_ROCM_FFT_CACHE_SIZE` (default 128, as on CUDA), the number of
+hipFFT plans the plan cache keeps; a smaller value only trades memory for plan
+rebuilds. It used to default to 8 to dodge a hang at the first plan eviction,
+which turned out to be a device-flag ordering bug in the backend, now fixed
+(lablup/mlxcel#1876). `MLX_ROCM_GPU_WATCHDOG_SECS` (default `0`,
 off) is the longest, in whole seconds, a single host wait for GPU work may
 block before it fails with an error (lablup/mlxcel#1804). A value that is not a
 non-negative integer (for example `10s`) is ignored with a stderr warning and

docs/installation.md

@@ -481,7 +481,7 @@ thinking model that spends its whole budget inside the thinking block reports
 | Memory estimation on UMA hosts | Correct. Measured on the tested configuration: the ROCm allocator reports a nonzero cap (76.80 GiB of the 96 GiB carve-out), which the estimator reads before it would ever reach host RAM, so nothing that fits the carve-out is refused for that reason (lablup/mlxcel#1805). |
 | Diagnostics | Report the AMD vendor, device name, `gfx` target and device memory, and no longer print a CUDA compute capability for it (lablup/mlxcel#1805). A binary whose compiled `gfx` list does not cover the device refuses to start rather than failing at the first kernel launch. |
 | `mlxcel-server` chat completions | Work for dense and affine MoE checkpoints, streaming and non-streaming; verified with `scripts/server_chat_smoke.sh`. |
-| Audio (speech to text, text to speech) | Works. The FFT primitive runs on hipFFT; `MLX_ROCM_FFT_CACHE_SIZE` bounds how many transform plans stay alive (lablup/mlxcel#1825). |
+| Audio (speech to text, text to speech) | Works. The FFT primitive runs on hipFFT; plans are cached up to `MLX_ROCM_FFT_CACHE_SIZE` (default 128, as on CUDA; lablup/mlxcel#1825, #1876). |
 | Windows, multiple GPUs, distributed inference | Not supported. |
 
 Decode throughput measured on the tested configuration, for orientation only

src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md

@@ -26,8 +26,9 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code
 15. **`SearchSorted` implemented.** The fork leaves `SearchSorted` a `NO_GPU` stub, and the exception crosses the cxx bridge into a `noexcept` extern, so it ends in `std::terminate`. mlxcel has no call site of its own; it is reached from MLX internals, because `residual_resample` in speculative decoding calls `fused_sample`, which falls back to `random::categorical` where no Gumbel-max kernel exists, and `categorical_impl` takes its inverse-CDF branch for 1-D logits, which is mlxcel's single-row sampling case. `mlx/backend/rocm/sort.hip` now carries the kernel, ported from `mlx/backend/cuda/sort.cu` and reusing this backend's `LessThan` so NaN ordering matches the CPU. Unlike the CUDA original it is grid-stride: the ROCm launch geometry caps the grid, so one index per thread would silently leave the tail of a large output unwritten. Verified against the CPU stream on both sides for duplicate runs, strictly increasing sequences, out-of-range queries, NaN, a single element, and 17M queries (past the 65535 x 256 grid extent). Applies to the fork; to be proposed there.
 16. **`Hadamard` implemented.** The fork leaves `Hadamard` a `NO_GPU` stub, which aborted every TurboQuant KV cache mode and `tests/wht_op.rs`. `mlx/backend/rocm/hadamard.hip` adds a single-stage radix-2 butterfly in shared memory for `m == 1` and a power-of-two last axis up to 8192, which is the whole of what mlxcel reaches: `wht()` asserts a power-of-two last axis and the KV cache advisor demotes Turbo modes for the models that would not have one, so the sizes seen here are head dims, chiefly 64, 128 and 256. Larger `n` and the `m` in {12, 20, 28} families throw a message naming the limit rather than returning a wrong answer; the CUDA path covers them by generating a radix-m codelet through NVRTC. The scale is the primitive's `scale_`, which MLX defaults to `1/sqrt(n)`, so the transform stays the involution mlxcel relies on; accumulation is in float regardless of element type. Like the CUDA original the grid is clamped at 65535 blocks and the kernel is grid-stride. Measured with `examples/wht_numeric_probe` against the M1 Ultra Metal reference: norm ratio 1.000000 to within 1e-5 at head_dim 64, 128 and 256, fp16 relative RMS about 2.0e-4 against Metal's 3.8e-4 to 5.0e-4, f32 relative RMS about 1e-7 matching Metal. The fp16 figure is lower than Metal's because this kernel keeps the butterfly in float in shared memory and rounds once on the way out, where Metal rounds between stages; more accurate than the reference is still inside the band, but it is worth knowing before someone reads the gap as a measurement error. Applies to the fork; to be proposed there.
 17. **Narrow index dtypes in `Gather` and `Scatter`.** Five dispatch sites in `mlx/backend/rocm/indexing.hip` picked the kernel's index type with `int32_t` for int32 and uint32 and `int64_t` for everything else, so int8, uint8, int16, uint16 and bool index arrays reached a kernel that read eight bytes per element from a one or two byte buffer. On ROCm that is an `HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION`: the GPU queue aborts and the process spins instead of failing. mlxcel reaches it with uint8, because the TurboQuant 4-bit codebook lookup is a `take` with uint8 indices, which is why every TurboQuant KV cache test faulted once the `Hadamard` stub was gone. Narrow index arrays are now widened to int32 once before the dispatch. The alternative, giving each of the eight integer dtypes its own instantiation, was implemented and measured first: it is also correct but takes this translation unit from 26 s to 348 s to compile, because the index type multiplies an existing product of twelve value dtypes and five to six index-count branches, so widening was kept at the cost of one copy pass on the paths that are broken today. Upstream does instantiate per index dtype, in both Metal (`backend/metal/indexing.cpp` puts the index dtype in the kernel name and builds the source from a template) and CUDA (`backend/cuda/indexing.cpp` passes it as an NVRTC template argument), but both compile the kernel at runtime and cache it, so the combinations cost nothing at build time. This backend compiles `indexing.hip` ahead of time, so the same strategy is paid by every build. The asymmetry is the compile model, not the strategy, and it is the argument for proposing widening rather than instantiation to the fork. Upstream already uses the same move for the offset type, choosing int32 or int64 by size rather than instantiating every combination. Verified against the CPU stream for int8, uint8, int16, uint16, int32, uint32, int64 and uint64 at four index shapes. Applies to the fork; to be proposed there.
-18. **`FFT` implemented through hipFFT.** The fork leaves `FFT` a `NO_GPU` stub, which aborted Kokoro text to speech and the Phi-4-multimodal audio front end, the only two mlxcel paths that reach a device FFT (the other audio front ends use the host-side transform in `src/audio/fft.rs`). `mlx/backend/rocm/fft.hip` ports `mlx/backend/cuda/fft.cu`, with two differences. This hipFFT has no `hipfftXtMakePlanMany` or `hipfftXtExec`, so the typed API is used instead: one `hipfftType` per plan and one exec entry point per transform, which is enough because MLX only asks for float32 and complex64. And hipFFT work cannot be a HIP graph kernel node, so execution goes through `CommandEncoder::launch_kernel`, the same path the rocBLAS and hipBLASLt GEMMs take. The plan cache defaults to 8 entries rather than the CUDA file's 128: with more hipFFT plans alive at once, creating the next one blocks forever inside the library with the GPU idle, reproduced deterministically at capacities 16, 24, 32 and 128 and never at 8, by a probe cycling 26 transform shapes. A plan outlives its cache entry while a pending execution still holds it, so the cap has to leave headroom. The root cause inside rocFFT is not identified; `MLX_ROCM_FFT_CACHE_SIZE` raises the cap for anyone measuring it. Verified against the CPU stream for rfft, irfft, round trips and complex fft at lengths 2, 8, 400, 512, 1024, 1200 and 2048 and batches up to 512, all within 1e-5 relative. Applies to the fork; to be proposed there.
+18. **`FFT` implemented through hipFFT.** The fork leaves `FFT` a `NO_GPU` stub, which aborted Kokoro text to speech and the Phi-4-multimodal audio front end, the only two mlxcel paths that reach a device FFT (the other audio front ends use the host-side transform in `src/audio/fft.rs`). `mlx/backend/rocm/fft.hip` ports `mlx/backend/cuda/fft.cu`, with two differences. This hipFFT has no `hipfftXtMakePlanMany` or `hipfftXtExec`, so the typed API is used instead: one `hipfftType` per plan and one exec entry point per transform, which is enough because MLX only asks for float32 and complex64. And hipFFT work cannot be a HIP graph kernel node, so execution goes through `CommandEncoder::launch_kernel`, the same path the rocBLAS and hipBLASLt GEMMs take. The plan cache holds 128 entries, as in the CUDA file (`MLX_ROCM_FFT_CACHE_SIZE` overrides it). It was 8 until lablup/mlxcel#1876: at 16 and above the first plan eviction hung inside `hipfftDestroy`, whose `hipFree` never returned. The cause was not rocFFT or the number of live plans but the order of `hipSetDeviceFlags` against the first HIP queue, fixed in entry 20. Verified against the CPU stream for rfft, irfft, round trips and complex fft at lengths 2, 8, 400, 512, 1024, 1200 and 2048 and batches up to 512, all within 1e-5 relative. Applies to the fork; to be proposed there.
 19. **`get_launch_args` removed.** The fork's `mlx/backend/rocm/kernel_utils.hpp` defined `get_launch_args` (a sized overload and an `array` overload) that fixed the block at 256 threads, ignored `shape`, `strides` and `large`, and capped the grid at 65535 blocks, returning that grid with nothing in its name or comment saying the kernel had to be grid-stride. A kernel written one index per thread with an early return past the end would leave every element past 65535 x 256 unwritten, with the output keeping whatever the allocation held (the failure shape of item 13). Upstream CUDA's helper of the same name (`mlx/backend/cuda/kernel_utils.cu` at the pin) does not cap x, takes a `max_block_dim`, and honours `large` through `get_2d_grid_dims`, so a kernel ported from the CUDA backend with its call intact would compile here and inherit the truncation. Nothing called it, in this overlay or in the fork at `75915908`; launch sites compute their own geometry. Most sites that clamp at 65535 blocks launch a grid-stride kernel (`binary.hip`, `unary.hip`, `hadamard.hip`, `sort.hip` among them), but not all: `SliceUpdate::eval_gpu`'s reduce-op path in `indexing.hip` clamps and launches `slice_update_op_kernel`, which does at most `NWORK` elements per thread and never reads `gridDim`, so an update larger than 65535 x 256 x `NWORK` elements leaves its tail unapplied. That is the same defect at a real call site, left for a separate fix because this item changes no kernel's geometry. Of the three options weighed in lablup/mlxcel#1874 (delete it, rename it to state the grid-stride requirement, or spill the overflow into another grid dimension the way upstream's `get_launch_args_general` spills y into z), deletion was taken: it cannot regress, and the other two keep a helper that only looks like upstream's. Spilling into y would not even make a one-index-per-thread kernel correct unless that kernel also reads `blockIdx.y`, and AMD caps each grid dimension at 2^32 - 1 threads (item 11), so a faithful port of upstream's geometry is not a drop-in either. In place of the definitions the header now declares `template <typename... Args> void get_launch_args(Args&&...) = delete;` with a comment giving this reason, so a ported call fails to compile at the explanation rather than as an undeclared name someone might "fix" by restoring the old helper. A fork sync that brings the old definitions back must drop them again. Applies to the fork; to be proposed there.
+20. **Blocking-sync flag set before the first HIP queue.** The fork set `hipSetDeviceFlags(hipDeviceScheduleBlockingSync)` in `rocm::device()`, but the allocator runs first: the first unified allocation's host read calls `hipStreamQuery(nullptr)` in `allocator::Buffer::raw_ptr()`, which creates the null stream's queue, and CLR gives that queue non-interrupt completion signals because active wait is still on. After the flag flips, CLR (`HwQueueTracker::ActiveSignal` in `rocvirtual.cpp`) no longer swaps such a signal for an interrupt one when a command needs a completion handler, and ROCr rejects async handlers on non-interrupt signals (`hsa_amd_signal_async_handler() failed to set the handler!` under `AMD_LOG_LEVEL=4`). The marker a device-wide sync enqueues on the null stream then never signals its waiter: `hipFree` hung in `std::condition_variable::wait` below `Device::SyncAllStreams`. rocFFT's plan destroy was the first thing to hit it (entry 18), but any `hipFree` or device-wide sync with pending null-stream work could. `rocm::ensure_device_flags()` in `device.cpp` now sets the flag once per device index, thread-safe and touching only that device, and runs before every path that can create a queue ahead of `rocm::device()`: `unified_malloc`, both branches of `raw_ptr()`, the async-pool free streams in the `RocmAllocator` constructor, `gpu::init()` and `staged_write`'s device sync; `rocm::device()` uses it too. If `hipSetDevice` fails it returns without recording the device, so a later call retries; if `hipSetDeviceFlags` fails it prints one line to stderr naming the device index and the HIP error and does not retry, so a regression of #1876 is visible. With it, `hipSetDeviceFlags` precedes the first `Created SWq` line in the `AMD_LOG_LEVEL=3` log, the 40-shape plan probe completes at plan-cache capacities 16, 32 and 128, and `tests/rocm_fft_plan_cache.rs` guards it (it times out with the ordering reverted). A standalone HIP reproducer without MLX or hipFFT (null-stream touch, then the flag, then `hipMalloc`, sync `hipMemcpy`, a null-stream kernel and `hipFree` in a loop) hangs at its first `hipFree` and completes without the early touch. Measured on gfx1151 with HIP 7.15.26333 and rocFFT 1.0.39 only. Applies to the fork; to be proposed there, and the reproducer is a candidate for a CLR report.
 
 ## Dropped from the vendored set
 

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp

@@ -130,6 +130,9 @@ inline void* unified_malloc(size_t size, bool& is_managed) {
   void* data = nullptr;
   hipError_t err;
   ensure_mlx_device_current();
+  // Often the first HIP work in the process, ahead of rocm::device(): put the
+  // device in blocking-sync mode before any queue exists on it (#1876).
+  ensure_current_device_flags();
 
   if (use_finegrained()) {
     err = hipExtMallocWithFlags(&data, size, hipDeviceMallocFinegrained);
@@ -326,6 +329,9 @@ RocmAllocator::RocmAllocator()
         }
       }
       (void)hipSetDevice(i);
+      // The free stream below is a queue on device i; flag the device first
+      // (#1876). This loop already switches to every pooled device.
+      ensure_device_flags(i);
       hipMemPool_t pool = nullptr;
       if (hipDeviceGetDefaultMemPool(&pool, i) == hipSuccess) {
         mem_pools_[i] = pool;
@@ -925,11 +931,17 @@ void* Buffer::raw_ptr() {
   auto& cbuf = *static_cast<rocm::RocmBuffer*>(ptr_);
 
   if (cbuf.device == -1) {
+    // The null-stream query creates the null stream's queue on first use, and
+    // a host read can come before rocm::device() ever runs. Flag the device
+    // the query is about to touch first, or that queue keeps completion
+    // signals a later hipFree waits on forever (#1876).
+    rocm::ensure_current_device_flags();
     if (hipStreamQuery(nullptr) != hipSuccess) {
       (void)hipStreamSynchronize(nullptr);
     }
   } else {
     // Discrete: host shadow (CUDA: move_to_unified_memory).
+    rocm::ensure_current_device_flags();
     (void)hipDeviceSynchronize();
     rocm::allocator().ensure_host_shadow(cbuf);
     return cbuf.host_shadow;