Repository navigation
build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802
Description
Part of #1801. Phase 0. Gates every other sub-issue.
Context
The feasibility spike in #1801 built the ROCm backend from NripeshN/mlx@75915908 on top of mlxcel's pin ml-explore/mlx@81ba1c6a, fixed six API-drift breaks, and validated it on gfx1151 (42/42 op checks, mlx-lm decode at fork parity). The result is a 121-file overlay: the fork's mlx/backend/rocm/ directory (106 files, already including the drift fixes and the mxfp4/mxfp8 scale-type dispatch fix) and 15 MLX core files merged at 81ba1c6a:
CMakeLists.txt
mlx/CMakeLists.txt
mlx/backend/common/buffer_cache.h
mlx/backend/common/compiled.cpp
mlx/backend/common/compiled.h
mlx/backend/gpu/primitives.cpp
mlx/compile.cpp
mlx/device.cpp
mlx/fast.cpp
mlx/fast.h
mlx/fast_primitives.h
mlx/io/safetensors.cpp
mlx/ops.cpp
mlx/primitives.cpp
mlx/stream.cpp
Today src/lib/mlx-cpp/CMakeLists.txt fetches MLX at the pin (line 122), applies whole-file overlays in mlx_apply_source_overlays (lines 18-82; patches-cuda/ only when MLX_BUILD_CUDA, lines 70-81), then add_subdirectory. src/lib/mlxcel-core/build.rs picks the backend by target_os plus the cuda feature, links CUDA in link_cuda() (line 492) and resolves architectures in detect_cuda_arch() (line 444). Nothing knows about ROCm.
Scope
Land the overlay and the build glue so that cargo build --features rocm produces a working GPU binary on a Linux AMD host. Nothing else in the runtime changes here; kernel routing, memory policy and quantization policy are Phase 1 sub-issues.
Implementation plan
- Overlay tree. The overlay is the first new member of mlxcelverse; keep it at
src/lib/mlx-cpp/patches-rocm/next topatches-cuda/so the later reorganization (refactor(mlx-cpp): organize mlxcel's MLX-side layer as mlxcelverse (backend overlays and mlxcel kernels) #1816) moves all backends together. Addsrc/lib/mlx-cpp/patches-rocm/mirroring the MLX tree with the 121 files. Addpatches-rocm/UPSTREAM(source repository, branch, commit75915908dfe5028335d318b10340313744fd3a8d, upstream base it was retargeted to, license MIT) and a shortpatches-rocm/README.mdthat states what the directory is, where it came from, and that files are whole-file overlays copied only for ROCm builds. - CMake. In
mlx_apply_source_overlays, copypatches-rocm/**(recursive,configure_file COPYONLY) only whenMLX_BUILD_ROCM. Fail configure with a clear message ifMLX_BUILD_ROCMandMLX_BUILD_CUDAare both on. ForwardMLX_BUILD_ROCM,CMAKE_HIP_ARCHITECTURESandMLX_ROCM_ARCHITECTURESto the MLX subproject. - Cargo features. Add
rocm = []tosrc/lib/mlxcel-core/Cargo.tomlandrocm = ["mlxcel-core/rocm"]to the rootCargo.toml. Enablingrocmtogether withcudaormetalis a build error. build_mlx. Whenrocmis on (Linux only), defineMLX_BUILD_ROCM=ON, keepMLX_BUILD_CUDA/METALoff, and pass the resolved HIP architectures.detect_rocm_arch(). Mirrorresolve_cuda_architectures()/detect_cuda_arch(): honorMLX_ROCM_ARCHITECTURESif set, else parserocminfo(Name: gfxNNNNof the GPU agents), else fail with a message telling the user to set the variable. Embed the compiled list in the binary the same way the CUDA arch list is, so a mismatch can be diagnosed at startup.link_rocm(). ResolveROCM_PATH(default/opt/rocm). Link the static kernel library produced by the ROCm backend plusamdhip64,rocblas,hiprand,hiprtcandhipblaslt(the set the backend's own CMake links). Emit-Wl,-rpath,$ROCM_PATH/lib: the spike host does not register ROCm in the loader path, and the binary must run withoutLD_LIBRARY_PATH.- Rebuild triggers. Add
cargo:rerun-if-changed=../mlx-cpp/patches-rocmnext to lines 241-242, andrerun-if-env-changedforROCM_PATHandMLX_ROCM_ARCHITECTURES. - Build cache. Confirm that a
rocmbuild and a default build use differentOUT_DIRs so the patched_deps/mlx-srcis never shared. If they can collide, include the backend in the cache marker checked bypurge_stale_mlx_cache(line 288). - Repository gates. Run
python3 scripts/ci/check_cross_repo_refs.pyand rewrite any bare#NNNreferences in vendored comments toNripeshN/mlx#NNNorml-explore/mlx#NNN. Runpython3 scripts/insert_apache_header.py --checkand confirm vendored files are skipped (all vendored.cpp/.hfiles carry an Apple copyright line;.hip/.hppare outside the gate). Runmake verify-kernel-dtype-keysif any vendored file containscuda_kernel(. Add aNOTICEentry for the ROCm backend (MIT, NripeshN/mlx contributors).
Acceptance criteria
-
cargo build --release --features rocmsucceeds on a Linuxgfx1151host and producesmlxcelandmlxcel-server. -
./target/release/mlxcel generate -m models/mlx/Qwen3-0.6B-4bit -p "Hello" -n 50 --temp 0runs on the GPU (not the CPU fallback) and produces coherent text, withLD_LIBRARY_PATHunset. - A second
cargo build --release --features rocmafter touching a bridge.cppreconfigures cleanly (the overlay copy is idempotent) and does not refetch MLX. - Editing a file under
patches-rocm/triggers a rebuild. -
--features rocm,cudaand--features rocm,metalfail with a clear error. - Default,
metalandcudabuilds compile no file frompatches-rocm/(verify from the configure log) and their existing gates stay green. -
check_cross_repo_refs.pyandinsert_apache_header.py --checkpass;NOTICEandpatches-rocm/UPSTREAMare present.
Validation
cargo build --release --features rocm
env -u LD_LIBRARY_PATH ./target/release/mlxcel generate -m models/mlx/Qwen3-0.6B-4bit -p "Hello" -n 50 --temp 0
readelf -d target/release/mlxcel | grep -E 'RPATH|RUNPATH|NEEDED'
touch src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp && cargo build --release --features rocm
cargo build --release --features rocm,cuda # must fail
Notes
- The spike needed
libopenblas-dev liblapack-dev liblapacke-devon the host (already an mlxcel Linux prerequisite). - mlxcel links
libmlx.astatically, so missingeval_gpusymbols fail at link time here, whereas the spike's shared Python build only failed at import.nm -D --undefined-only libmlx.so | c++filt | grep mlx::corewas how the spike found them. - Kernel routing (refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803) is not in scope. Qwen3-0.6B works without it because the affected fused paths have graph fallbacks; MoE and BitNet may not.
References
- Overlay mechanism:
src/lib/mlx-cpp/CMakeLists.txt:18-82, pin:122 - Build glue:
src/lib/mlxcel-core/build.rs(rerun-if-changed240-242,purge_stale_mlx_cache288,build_mlx343,detect_cuda_arch444,link_cuda492) - Features:
Cargo.toml:69-73,src/lib/mlxcel-core/Cargo.toml:12-16 - ROCm backend link set:
mlx/backend/rocm/CMakeLists.txtin the fork (target_link_libraries(mlx ...))
Activity
[-]build(rocm): vendor the xcelverse overlay and add the rocm cargo feature[/-][+]build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature[/+]fix(rocm): fix mxfp gather_qmm dispatch and arg_reduce launch grid
fix(rocm): prune stale ROCm overlay files, require the overlay
feat(rocm): add the ROCm backend via the mlxcelverse overlay (#1818)
test(rocm): add the gfx1151 correctness matrix and its gate
refactor(core): route custom kernels by GPU backend kind
refactor(core): route custom kernels by GPU backend kind
ci(rocm): build, link and generate on a self-hosted gfx1151 runner
ci(rocm): build, link and generate on a self-hosted gfx1151 runner
ci(rocm): build, link and generate on a self-hosted gfx1151 runner (#…
test(rocm): add a local verify-rocm gate with a shared build-and-gene…