feat(rocm): implement FFT, Hadamard and SearchSorted on ROCm (audio, TurboQuant KV cache, stochastic speculative decoding abort today) #1825

Part of #1801. Phase 1. Found by the ROCm test gate in #1809.

Context

The ROCm backend vendored in mlxcelverse leaves three MLX primitives as NO_GPU stubs (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp:32,33,45) that throw "<name> has no ROCm implementation." The exception crosses the cxx bridge into a noexcept extern, so every hit ends in std::terminate.

FFT scope is narrow; an earlier revision of this issue overstated it. src/audio/fft.rs is a bounded host-side FFT (Bluestein for non-power-of-two lengths, radix-2 otherwise, capped at MAX_REAL_FFT_LEN = 1 << 17), and the Whisper, Gemma 4, Gemma3n, Nemotron-Omni and Inkling front ends use it rather than a device FFT. The ROCm gate agrees: the only tests that hit the FFT stub are audio::phi4mm::feature_extractor::tests::deterministic_waveform_matches_pinned_speechlib_features, audio::phi4mm::feature_extractor::tests::pinned_processor_relabels_24khz_before_feature_extraction and models::kokoro::stft::tests::stft_istft_round_trip_is_unit_scale. Speech to text in general is not blocked.

Hadamard carries a normalization contract that a port can get wrong while still passing a naive test. mlxcel calls hadamard_transform with no scale argument, so MLX applies the orthonormal 1/sqrt(N) scaling (mlx/ops.cpp, and src/lib/mlxcel-core/src/ops.rs documents the same), which makes the transform an involution: wht(wht(x)) must return x. A ROCm kernel written with a scale of 1 or 1/N would leave a round-trip test that only checks shape or finiteness passing while silently corrupting KV values, so the involution is a required acceptance check below.

The Hadamard sizes that matter are small powers of two. src/lib/mlxcel-core/src/ops.rs asserts the last axis is a power of two, and src/execution/kv_cache_advisor.rs:395 demotes Turbo modes to int8 for models whose head_dim is not a power of two (Phi-2's 80, MLA low-rank latents). MLX's general constraint is n = m * 2^k for m in {1, 12, 20, 28}, but the m in {12, 20, 28} families are unreachable from mlxcel and use a different normalization. The port therefore needs head_dim-sized powers of two on the last axis, chiefly 64, 128 and 256, with fp16 as the default dtype and fp32 also exercised (the K rotation path casts before the transform, cache/turbo/quant.rs:739-741).

SearchSorted is reached indirectly. Nothing in the Rust code or the cxx bridge calls MLX's searchsorted; TurboQuant's nearest-centroid lookup is a host-side binary search (src/lib/mlxcel-core/src/cache/turbo/codebook.rs:271). The aborts come from MLX internals on the sampling path: residual_resample (src/lib/mlxcel-core/src/speculative/stochastic_accept.rs:576) calls ffi::fused_sample, which falls back to the categorical graph path on a backend with no Gumbel-max kernel, and categorical_impl in mlx/random.cpp takes its inverse-CDF branch (categorical_inverse_cdf, which calls searchsorted on a 1-D CDF with side = "right") whenever the logits are 1-D, which is mlxcel's single-row sampling case. The other branch is gumbel + argmax and needs no searchsorted. Two consequences: the required kernel is just a binary search over a 1-D sorted CDF, small and cheap to port; and landing the ROCm Gumbel-max kernel under #1803 would remove this dependency from the fused sampling path, though the graph fallback would still need it. The 13 aborting tests are in speculative::distribution_tests and speculative::stochastic_accept.

Scope

Give the ROCm backend real implementations of the three primitives, inside src/lib/mlx-cpp/patches-rocm/, so the listed paths run on the GPU. Where a GPU kernel is not ready, a correct fallback that does not abort is acceptable as an intermediate step.

Implementation plan

1. FFT: port mlx/backend/cuda/fft.cu to hipFFT, including the plan cache, real and complex transforms, and the inverse normalization MLX expects. Link hipfft in link_rocm() (src/lib/mlxcel-core/build.rs:537) and in the backend CMake; the host already ships hipfft and rocfft CMake packages.
2. Hadamard: port mlx/backend/cuda/hadamard.cu with the 1/sqrt(N) scale, prioritising the small power-of-two path over the large-n two-stage path.
3. SearchSorted: port the searchsorted implementation from mlx/backend/cuda/sort.cu, covering both the left and right sides.
4. Record each change in src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md.

Acceptance criteria

• wht(wht(x)) round-trips to x within float tolerance at head_dim 64, 128 and 256 in fp16 on ROCm, and the TurboQuant cache tests and tests/wht_op.rs pass.
• rfft and irfft match the CPU within float tolerance, and the Kokoro STFT round-trip test and the Phi-4-multimodal feature extractor tests pass.
• searchsorted matches the CPU on both the left and right sides, and the speculative::distribution_tests and speculative::stochastic_accept tests pass.
• A Kokoro synthesis and a TurboQuant KV cache generation both run on ROCm.

References

• ROCm stubs: src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp (NO_GPU(FFT), NO_GPU(Hadamard), NO_GPU(SearchSorted))
• Gate results: test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809
• Backend-kind kernel routing and the ROCm Gumbel-max path: refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803

Provenance

The scope corrections above come from cross-checking the Metal side of each primitive against the ROCm gate log on gfx1151; the earlier claim that all audio front ends are blocked was wrong.