fix(rocm): hipFFT blocks when too many plans are alive, root cause unknown

[inureyes]

Part of #1801

Problem / Background

The ROCm FFT port caps its hipFFT plan cache at 8 entries, against the CUDA file's 128, because past roughly a dozen live plans the next hipfftMakePlanMany64 never returns. The GPU is idle and the calling thread sleeps on a pipe read. The cap bounds the symptom; the cause is not identified, so 8 is an empirical limit rather than a considered design choice, and every FFT-heavy workload (Kokoro text to speech, the Phi-4-multimodal audio front end) pays repeated plan rebuilds for it.

Current Behavior

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip:128-141 sets default_capacity to 8 on the LRUBytesKeyCache, with a comment recording the measurement. Upstream's mlx/backend/cuda/fft.cu:90-92 uses 128 at the same call site. LOCAL_FIXES.md entry 18 carries the same record.

Measured on a Radeon 8060S (gfx1151), ROCm 10.0.0, with a probe cycling 26 transform shapes (rfft, irfft, round trips, and complex fft at lengths 2, 8, 400, 512, 1024, 1200, 2048, batches 3, 4, 512): capacity 8 completes every run, three runs in a row; capacities 16, 24, 32, and 128 all hang, and always at the same position in the sequence rather than on a particular shape. Reordering the cases moves the hang to whatever runs last, and 60 repeats of a single shape never hang, which is what isolates it to the number of live plans rather than the number of executions.

It is not a plain plan limit in hipFFT. A standalone program outside MLX creates 64 plans without trouble, both with hipFFT's own work-area allocation and with hipfftSetAutoAllocation(handle, 0) plus a manual work area and an execution per plan. So it is an interaction with how MLX holds or executes them. One relevant detail: a plan outlives its cache entry while a pending execution still holds its shared_ptr (the cache stores std::shared_ptr<HipFFTPlan>, fft.hip:138), so the live count can exceed the configured capacity, which is why the cap has to leave headroom. Note also that this port sets hipfftSetAutoAllocation(handle, 0) at fft.hip:180 and allocates the work area per execution from MLX's pool, which upstream's CUDA path does not do in the same way.

MLX_ROCM_FFT_CACHE_SIZE raises the cap for anyone measuring it.

Scope

In scope: investigating the block, and src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip plus LOCAL_FIXES.md entry 18 if the cap turns out to be avoidable.

Out of scope: changing the transform math or the typed-API port itself, which is verified correct to 1e-5 relative against the CPU stream. Changing LRUBytesKeyCache.

Proposed Solution

Find the cause and, if it is ours, remove the cap. Starting points, in order:

1. Capture where the blocked thread actually is: attach with gdb -p and take a backtrace of all threads. The pipe read suggests a runtime waiting on something rather than a GPU stall, so the frame above it is the answer.
2. Check whether rocFFT's runtime kernel compilation is involved, and whether its cache path or a lock is the thing being waited on. ROCFFT_RTC_CACHE_PATH and the rocFFT RTC log knobs are the levers.
3. Re-run the probe with hipfftSetAutoAllocation left on, so hipFFT owns the work area, to see whether the manual work area at fft.hip:180 is implicated.
4. Check whether releasing plans eagerly on eviction, rather than when the last shared_ptr drops, changes the threshold. That separates "too many handles alive" from "too many handles in the cache".

If the limit is genuinely inside rocFFT, stop at a minimal standalone reproducer and report it.

Implementation Notes

• Reuse: the existing 26-shape probe and MLX_ROCM_FFT_CACHE_SIZE are the measurement apparatus; do not write a new harness. examples/wht_numeric_probe.rs is the in-tree pattern for a probe binary if one needs to be checked in.
• Constraints: gfx1151 and ROCm 10.0.0 are the only measured configuration. Any conclusion should say whether it was checked on another GPU or ROCm version. The hang is a hard hang with no timeout, so every experiment needs an external kill.
• Edge cases: the headroom point above means a raised cap is only safe if the live-plan count, not the cache size, is what is bounded. Measure the live count, not the configured capacity.
• Error handling: if the cap stays, check_hipfft_error cannot help, because the call never returns. Any mitigation has to be a bound on plan count rather than an error path.

Acceptance Criteria

• The cause is identified and written down, or, if it is not found, the dead ends are recorded (what was checked, what it ruled out) so the next person does not repeat them.
• If the block turns out to be avoidable on our side, the cap is raised or removed and the 26-shape probe passes three consecutive runs at the new capacity on gfx1151.
• If the limit is genuinely in rocFFT, a minimal standalone reproducer exists, the report is filed upstream, and the comment at fft.hip:128 cites that report by URL instead of saying the cause is unidentified.
• LOCAL_FIXES.md entry 18 is updated to match whatever the outcome is.
• FFT numerical results are unchanged: rfft, irfft, round trips, and complex fft within 1e-5 relative against the CPU stream at the lengths and batches already covered.

Verification

MLX_ROCM_FFT_CACHE_SIZE=128 cargo test --release --features rocm fft

A pass is three consecutive completions of the 26-shape probe at the raised capacity with no hang, plus the existing FFT correctness checks still within 1e-5 relative. If the cap stays at 8, a pass is instead the written-down cause or dead ends plus the upstream report link in the source comment.

Technical Considerations

Related: #1825 (the FFT implementation this came out of), PR #1861 (the port), #1813 (upstreaming the fork-side local fixes). The hang is a hard hang, so this should not be worked on a machine anyone else is using.