Skip to content

test(rocm): promote the FFT probe to a reusable example - #1878

Open
inureyes wants to merge 1 commit into
mainfrom
feature/rocm-fft-numeric-probe
Open

inureyes wants to merge 1 commit into
mainfrom
feature/rocm-fft-numeric-probe

Conversation

@inureyes

Copy link
Copy Markdown
Member

The FFT work in #1825 and the plan-cache hang in #1876 were both established with throwaway C++ programs that linked libmlx directly and lived only in a session scratchpad. #1876 asks whoever picks it up to re-run that stress with hipfftSetAutoAllocation left on, and there was nothing to re-run: the issue records the measurements as one person's, not as something a reader can reproduce.

examples/fft_numeric_probe replaces both, next to wht_numeric_probe and built the same way, so it needs no separate toolchain. Two modes, because the questions differ and one of them hangs.

The default compares rfft, irfft, a round trip and a complex fft against the CPU stream, at the lengths mlxcel asks for (512 for the Phi-4-multimodal front end, 1200 for Kokoro's STFT), a non-power-of-two, and batches from 3 to 512. The CPU is the reference rather than another host because two accelerators are not expected to agree bit for bit, and the CPU arm exists wherever this runs, so a disagreement here belongs to this backend.

--plans N creates N distinct transform shapes in one process, which is what makes a plan cache hold N plans at once. That is the #1876 reproduction: on gfx1151 it completes 40 plans at the default MLX_ROCM_FFT_CACHE_SIZE of 8 and stops at the 17th with the variable set to 32, with no error and the GPU idle. Run it under timeout when scripting it, since the failure is a hang.

Measured on gfx1151 with --features rocm: every correctness comparison is within 5e-7 relative, against a 1e-5 threshold. The example also builds and runs without a GPU backend, where both arms resolve to the CPU and the comparison is degenerate but the program is still exercised.

Refs #1825, #1876

Validated on gfx1151: cargo fmt --all -- --check, scripts/insert_apache_header.py --check, scripts/ci/check_cross_repo_refs.py, the correctness mode under --features rocm, the same mode on a default (CPU) build, and both --plans 40 runs described above.

The FFT work in #1825 and the plan-cache hang in #1876 were both established with throwaway C++ programs that linked libmlx directly and lived only in a session scratchpad. #1876 asks whoever picks it up to re-run that stress with `hipfftSetAutoAllocation` left on, and there was nothing to re-run: the issue records the measurements as one person's, not as something a reader can reproduce.

`examples/fft_numeric_probe` replaces both, next to `wht_numeric_probe` and built the same way, so it needs no separate toolchain. Two modes, because the questions differ and one of them hangs.

The default compares `rfft`, `irfft`, a round trip and a complex `fft` against the CPU stream, at the lengths mlxcel asks for (512 for the Phi-4-multimodal front end, 1200 for Kokoro's STFT), a non-power-of-two, and batches from 3 to 512. The CPU is the reference rather than another host because two accelerators are not expected to agree bit for bit, and the CPU arm exists wherever this runs, so a disagreement here belongs to this backend.

`--plans N` creates N distinct transform shapes in one process, which is what makes a plan cache hold N plans at once. That is the #1876 reproduction: on gfx1151 it completes 40 plans at the default `MLX_ROCM_FFT_CACHE_SIZE` of 8 and stops at the 17th with the variable set to 32, with no error and the GPU idle. Run it under `timeout` when scripting it, since the failure is a hang.

Measured on gfx1151 with `--features rocm`: every correctness comparison is within 5e-7 relative, against a 1e-5 threshold. The example also builds and runs without a GPU backend, where both arms resolve to the CPU and the comparison is degenerate but the program is still exercised.

Refs #1825, #1876
@inureyes inureyes added type:test Test related changes priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific status:review Under review labels Sep 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific priority:medium Medium priority status:review Under review type:test Test related changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant