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

This branch has not been deployed

No deployments
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