fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 1. Depends on #1802.
## Context
Upstream ml-explore/mlx#3742 made `array::is_available()` throw on a failed launch and consume the error, carried through `Event::error()` (`std::atomic`). The mlxcel bridge relies on this: the drain in `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp` (around lines 5571-5610) assumes a failed GPU step surfaces as an exception that becomes a Rust error.
The ROCm overlay in mlxcelverse only has storage for the error (added during the retarget in #1801); nothing in the ROCm command encoder, event or completion path ever calls `set_error`. During the spike, failures looked like this instead:
- mxfp4 `quantized_matmul` hung until killed;
- the unfixed mxfp8 path produced NaN and then `HSA_STATUS_ERROR_MEMORY_FAULT` with a queue hang dump;
- a failed `hipLaunchKernel` ("invalid configuration argument") threw from the launch site, but asynchronous faults did not.
A user-facing inference server cannot tell these apart from a slow model.
## Scope
Make ROCm launch failures and asynchronous GPU faults reach `Event::error` so that MLX throws and mlxcel reports an error. Deadlock detection for kernels that spin forever is out of scope beyond a documented watchdog option.
## Implementation plan
1. Audit the ROCm command encoder (`mlx/backend/rocm/device.*`, `worker.*`, `event.hip`, `eval.cpp`) for every `hip*` call whose `hipError_t` is ignored (the build emits `-Wunused-value` warnings for these, e.g. `hipGraphDestroy` in `device.h`). Check them and convert failures into an MLX `Error`.
2. On stream completion, query `hipStreamQuery`/`hipGetLastError` (or the completion callback status) and call `set_error` on the signaled event, mirroring the Metal completion handler (`mlx/backend/metal/device.cpp`, `event.set_error(error_)`) and the scheduler path (`mlx/scheduler.cpp`).
3. Treat an HSA memory fault that kills the queue as fatal for the stream: mark pending events with an error so waiters return instead of blocking forever.
4. Optional: an `MLXCEL_GPU_WATCHDOG_SECS` style timeout in the bridge drain that reports a stuck GPU step as an error with the last kernel name.
5. Keep the change inside `patches-rocm/` and record it in `patches-rocm/UPSTREAM` as a local fix to upstream later (#1813).
## Acceptance criteria
- [ ] A deliberately invalid launch (for example a kernel launched with an oversized block) returns an error to Rust through the bridge instead of aborting.
- [ ] An asynchronous device fault (for example an out-of-bounds read in a test kernel) makes the waiting `eval` throw within a bounded time instead of hanging.
- [ ] `mlxcel-server` returns an HTTP error for a request whose forward pass fails, and keeps serving later requests if the device is still usable.
- [ ] No behavior change on Metal or CUDA.
## Validation
A small test binary or `#[test]` behind `--features rocm` that triggers each failure class, plus a manual run of the pre-fix mxfp4 repro from #1808 showing an error instead of a hang.
## References
- Upstream contract: ml-explore/mlx#3742
- Bridge drain: `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp` (around 5571-5610)
- ROCm event storage added in the retarget: `mlx/backend/rocm/event.hip` (`EventImpl::error`)
Contributor guide
Research direction
Start with the ROCm files named in the issue—mlx/backend/rocm/device.*, worker.*, event.hip, and eval.cpp—and compare their completion handling with mlx/backend/metal/device.cpp and mlx/scheduler.cpp. Read the bridge drain in src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp around lines 5571-5610, then run the planned ROCm failure tests and mxfp4 repro. Done means launch and asynchronous faults reach Rust as errors, HTTP requests fail without hanging, and Metal/CUDA behavior is unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100