Repository navigation
test: add reliable FFT numeric and plan-cache probes - #1878
Merged
Merged
Conversation
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
added
type:test
inureyes
added a commit
that referenced
this pull request
Sep 29, 2026
…#2049) ## Summary The ROCm hipFFT plan cache was capped at 8 because the first plan eviction hung. The hang was not in rocFFT and did not depend on the number of live plans. The backend set `hipSetDeviceFlags(hipDeviceScheduleBlockingSync)` too late, after the allocator had already created the null stream's HIP queue. This PR sets the flag before the first queue exists and raises the cache back to the CUDA value of 128. ## Root cause - The probe (`fft_numeric_probe --plans 40`, from PR #1878) hung at cap 16 after `plan 7` and at cap 32 after `plan 15`. Each iteration creates two plans, so the hang was always at the first eviction, meaning the first `hipfftDestroy`. - rocgdb on the hung thread: `hipfftDestroy -> rocfft_plan_destroy -> ~ExecPlan -> ~TreeNode -> gpubuf_t::free -> hipFree -> ... -> std::condition_variable::wait`. The GPU was idle and no RTC helper process was running. - The last line under `AMD_LOG_LEVEL=4` was `rocvirtual.cpp:844 hsa_amd_signal_async_handler() failed to set the handler!`. The rejected signal had no event mailbox, so it was a non-interrupt signal, and ROCr refuses async handlers on those. - CLR created that signal (`HSA_AMD_SIGNAL_AMD_GPU_ONLY`) for the null stream's queue when `allocator::Buffer::raw_ptr()` called `hipStreamQuery(nullptr)`. The flag came afterwards from `rocm::device()`. Once active wait is off, `HwQueueTracker::ActiveSignal` stops swapping non-interrupt ring signals for interrupt ones. The marker that `hipFree`'s `SyncAllStreams` enqueues on the null stream lands on one of these stale signals, its handler is rejected, and the wait never returns. - This is not specific to FFT. Any `hipFree` or device-wide sync with pending null-stream work could hang the same way. rocFFT was simply the first code to do null-stream work plus `hipFree` at runtime. The fix covers all of these cases. - Not characterized: why the threshold falls at the first eviction for caps of 16 or more but not at 12 or below. It probably depends on the position in the null stream's signal ring. Standalone reproducer without MLX or hipFFT (`hipcc --offload-arch=gfx1151 -O2`): ```cpp #include <hip/hip_runtime.h> #include <cstdio> #include <cstdlib> #include <vector> __global__ void k(float* p, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) p[i] += 1.0f; } int main(int argc, char** argv) { bool early_touch = argc > 1 && argv[1][0] == '1'; int iters = argc > 2 ? atoi(argv[2]) : 200; if (early_touch) (void)hipStreamQuery(nullptr); // creates the null-stream queue now (void)hipSetDeviceFlags(hipDeviceScheduleBlockingSync); // flips active wait off hipStream_t s; hipStreamCreateWithFlags(&s, hipStreamNonBlocking); std::vector<float> h(1024, 1.0f); std::vector<void*> live; for (int i = 0; i < iters; ++i) { float* p; hipMalloc(&p, 4096); hipMemcpy(p, h.data(), 4096, hipMemcpyHostToDevice); hipLaunchKernelGGL(k, dim3(4), dim3(256), 0, nullptr, p, 1024); live.push_back(p); if (live.size() > 8) { hipFree(live.front()); live.erase(live.begin()); } printf("iter %d\n", i); fflush(stdout); } printf("done\n"); } ``` `timeout 30 ./repro 1 200` hangs at the first `hipFree` (the last line printed is `iter 7`, exit 124). `./repro 0 200`, without the early touch, prints `done`. The reproducer is ready for an upstream CLR report as a follow-up. None has been filed. ## Changes - `rocm::ensure_device_flags(int)` / `ensure_current_device_flags()` in `device.cpp` set BlockingSync once per device index. They are thread-safe, touch only the requested device, and restore the caller's current device. `rocm::device()` now uses them. They also run before every path that can create a queue ahead of `rocm::device()`: `unified_malloc`, both branches of `raw_ptr()`, the opt-in async-pool free streams, `gpu::init()`, and the device sync in `staged_write`. The flag value itself is unchanged. If `hipSetDevice` fails the helper returns without recording the device, so a later call retries; if `hipSetDeviceFlags` fails it prints one line to stderr naming the device and the HIP error, so a return of this hang would not be silent. - `fft.hip`: the plan cache is back to 128 and the comment states the real cause. The comment in `execute_fft` claimed that the captured plan outlives the call. In fact `launch_kernel` runs the lambda synchronously and drops it, and eviction is safe for the eager launch because `hipFree` synchronizes first. - `tests/rocm_fft_plan_cache.rs`: a new ROCm integration test. A child process runs with `MLX_ROCM_FFT_CACHE_SIZE=16` and a 120 s budget. It does 48 rfft/irfft round trips at distinct lengths, then GPU-vs-CPU checks within 1e-5 for rfft, irfft, round trip and complex fft at lengths 2, 8, 400, 512, 1024, 1200 and 2048, plus a 512-batch rfft. The test is picked up by `verify-test-rocm` like every other file in `tests/`. - `LOCAL_FIXES.md` entry 18 is updated and entry 20 is new; main was merged into the branch to pick up entry 19 from PR #2046. The docs no longer describe the old cap of 8. ## Verification Run on gfx1151 (Radeon 8060S) with HIP 7.15.26333 (ROCm 10.0 core) and rocFFT 1.0.39. No other GPU or ROCm version was checked. - `MLX_ROCM_FFT_CACHE_SIZE=128 timeout 120 target/release/examples/fft_numeric_probe --plans 40`, three consecutive runs: all three exited 0 with `created 40 plans without blocking` and 40/40 round trips ok. At cap 128 this run evicts nothing (80 plans), so the eviction cases below are the ones that matter. - The same probe at cap 16 and at cap 32, which hung before the fix after `plan 7` and `plan 15`: both exited 0 with 40/40 ok. Cap 16 with `--plans 100` also completed. - Probe correctness mode (`timeout 120 .../fft_numeric_probe`): all 25 comparisons are within 1e-5 of the CPU stream. The worst is 4.2e-7 (round trip at n=2048 and n=1200). - Ordering proof: `AMD_LOG_LEVEL=3 MLX_ROCM_FFT_CACHE_SIZE=16 timeout 120 .../fft_numeric_probe --plans 4` logs `hipSetDeviceFlags ( 4 )` on line 26, then `hipExtMallocWithFlags` on line 32 and `hipStreamQuery ( <null> )` on line 50, and the first `Created SWq` on line 52. Before the fix, `hipSetDeviceFlags` came after the null-stream `Created SWq`. - `cargo test --features rocm --test rocm_fft_plan_cache`: passed twice (9.97 s and 6.91 s). With the early `ensure_*_device_flags` calls commented out (the original order, where only `rocm::device()` sets the flag), it failed: `the child did not finish within 120s`, after 8 of 48 round trips, which is the first eviction. - `cargo clippy --features rocm --test rocm_fft_plan_cache -- -D warnings`, `cargo test --features rocm --test dead_doc_pointers`, and `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt`: all pass. - Re-run by the orchestrator on the final head after merging main: `cargo test --features rocm --test rocm_fft_plan_cache --test dead_doc_pointers` pass (child finishes in about 7 s); the same regression test with the early calls commented out fails with `the child did not finish within 120s`; the probe at cap 128 three consecutive times, at caps 16 and 32, and in correctness mode all complete; the fast gates pass; `scripts/ci/rocm_smoke.sh` with `MLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bit` passes (32 tokens generated). - Standalone reproducer: as described above (exit 124 at `iter 7` with the early touch, `done` without it). ## Not verified Metal and CUDA are not available on this host. This change touches only `patches-rocm/` (ROCm-only) and a `#![cfg(feature = "rocm")]` test, so the Metal and CUDA paths are unchanged. The opt-in `MLX_GRAPH_PREFILL_REPLAY=1` capture path for FFT was not examined. ## Changes during review Review and security passes found no CRITICAL or HIGH issues. Follow-up commits merged main (LOCAL_FIXES entry 19 from PR #2046), numbered the entries in the `fft.hip` comment, made the helper's failure paths visible as described above, and made the test kill its child if polling it fails. Known pre-existing issues outside this diff, not changed here: `lru_cache.h` evicts from an empty list when `MLX_ROCM_FFT_CACHE_SIZE=0`, and `std::stoul` throws on a non-numeric value. Refs #1825, #1878, #2033. Closes #1876
Merge current main and make both numerical modes return failure for incorrect or non-finite results. Reject invalid and overflowing plan counts before device work, add odd-length FFT coverage, and clarify backend and timeout usage. Four unit tests, 35 CUDA/CPU comparisons, 40 stress shapes, CLI exit-code checks, and full CUDA workspace/all-targets Clippy pass. Preserve historical ROCm evidence without claiming a fresh AMD run. Refs #1825, #1876
inureyes
changed the title
Member
Author
|
Final integration verification is complete on main
GitHub currently reports successful PR-head check rollups for #2304, #2138, and #2296; the other five rollups remain pending. These pending runs are not counted as passes above. The final-main Metal workflow and local CUDA suite are complete. Physical ROCm execution and M5-specific validation remain unverified; #2158 and #1986 remain open and outside this completed batch. The local main checkout is clean and matches the verified remote main. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replace the throwaway FFT programs from #1825 and #1876 with a reusable mlxcel-core example. Default mode compares rfft, irfft, round trips, and complex FFT with the CPU stream.
--plans Nexercises distinct forward/inverse shapes and cache eviction; retained plan count depends on the backend limit.Integrate main
e72b2b7aand MLX243bdf09; the latest README-only advance preserves the validated production source. Fix review findings: numerical failures previously printed FAIL but exited successfully, malformed counts silently defaulted to 40, and large counts could overflow dimensions. Both numerical modes now preserve failures and exit 1; invalid, zero, extra, and out-of-range arguments exit 2. Add odd lengths 3, 7, and 401. Document explicit backend features, degenerate CPU-only smoke tests, and building before applying the runtime timeout.Validation on NVIDIA GB10:
No blocking implementation, security, or performance finding remains. Bilingual reports preserve source and binary hashes. Earlier gfx1151 measurements (below 5e-7, cache size 8 completing 40 shapes, size 32 hanging near shape 17) remain historical; AMD hardware is unavailable here, and this PR does not claim to fix the hipFFT hang or allocation policy.
Refs #1825, #1876