epic: AMD GPU (ROCm) backend on Linux via mlxcelverse

[inureyes]

Summary

Program of work to run mlxcel model inference with GPU acceleration on Linux hosts with AMD GPUs (ROCm/HIP). The first target is RDNA 3.5 (gfx1151, Strix Halo), where a feasibility spike already produced correct results and usable 4-bit decode speed.

The approach is to vendor an existing ROCm backend for MLX into mlxcel as a set of source overlays and apply it on top of the same pinned ml-explore/mlx commit that the Metal and CUDA builds use. It becomes the ROCm part of mlxcelverse, the name for mlxcel's MLX-side layer (see Design decisions). The ROCm overlay is copied into the MLX tree only when building with a new rocm cargo feature, so Apple Silicon and CUDA builds are untouched.

The ROCm backend comes from the rocm-support branch of NripeshN/mlx (MIT), which is the head of the upstream draft ml-explore/mlx#2300 and the only actively developed ROCm line for MLX today (tracking request: ml-explore/mlx#2556). Downstream projects such as lemonade-sdk/lemon-mlx-engine already build against it.

Why an in-repo overlay and not a separate MLX fork repository

The fork was measured against its upstream merge base (39886de4 to 75915908):

Part	Size	Nature
mlx/backend/rocm/	106 new files, +43,904 lines	Purely additive. Upstream never touches this directory.
MLX core glue	17 files, about +345/-65 lines	Mostly #ifdef MLX_USE_ROCM hooks (default device, fast::hip_kernel declaration, custom-kernel stubs, CMake option) plus a few general tweaks (buffer donation, buffer cache, a quantized-matmul fallback).
Everything else	Python bindings, tests, docs, bench scripts	Not needed by mlxcel.

mlxcel's C++ bridge is written against the API of its pinned MLX commit, so any ROCm source has to follow that pin no matter where it lives. A separate repository would add a second pin, an ancestry rule between the two, and changes to both pin parsers (they require ml-explore/mlx in the repository URL), while the per-bump work stays the same. Vendoring keeps one pin and lets a pin bump fix the ROCm side in the same PR, using the overlay discipline mlxcel already applies to its Metal and CUDA patches.

This is different from porting MLX's CUDA backend to HIP by hand inside the overlays, which would mean replacing CUTLASS, CCCL and cudnn-frontend (10k+ lines). That option is not pursued.

Design decisions

0. mlxcelverse. mlxcelverse names everything mlxcel builds on top of upstream MLX, in two kinds: (a) per-backend source overlays that replace or add MLX files (today src/lib/mlx-cpp/patches/ for Metal and CUDA and src/lib/mlx-cpp/patches-cuda/; this epic adds patches-rocm/), maintained by a 3-way merge and a line-by-line review on every pin bump; and (b) mlxcel's own kernels, fusions and extension functions built on MLX's public custom-kernel APIs (fast::metal_kernel, fast::cuda_kernel, and on ROCm fast::hip_kernel), today in src/lib/mlx-cpp/turbo/ and the bridge, which only need public-API compatibility across bumps. This epic stays scoped to the ROCm members: the overlay (build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802) and ROCm kernel ports (perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814). Reorganizing the existing tree under the mlxcelverse name, with no change to any build output, is tracked separately in refactor(mlx-cpp): organize mlxcel's MLX-side layer as mlxcelverse (backend overlays and mlxcel kernels) #1816.
1. Layout. src/lib/mlx-cpp/patches-rocm/ mirrors the MLX tree: the whole mlx/backend/rocm/ directory plus 15 core files, all as whole-file overlays (no diff patches). Whole-file configure_file COPYONLY is idempotent across the reconfigures that build.rs triggers, which diff patches are not.
2. ROCm-only copy. The overlay is copied only when MLX_BUILD_ROCM is on, following the patches-cuda/ precedent, so Metal and CUDA builds never compile a ROCm-modified core file. MLX_BUILD_ROCM together with MLX_BUILD_CUDA is rejected at configure time.
3. Enable by cargo feature. --features rocm (root) forwards to mlxcel-core/rocm. Feature combinations get separate build-script output directories, so a ROCm-patched _deps/mlx-src is never reused by another build.
4. Single MLX pin. The ROCm build fetches the same ml-explore/mlx commit as every other build. The pin parsers do not change.
5. Provenance. patches-rocm/UPSTREAM records the source repository, branch, commit and license; NOTICE gets an entry; vendored files keep their original headers and never receive a Lablup header.
6. Quantization modes ROCm cannot run are converted where possible. Load-time policy converts unsupported modes to affine (for example NVFP4 to affine 4-bit, following the existing dense repack path) and rejects with a clear message only when conversion is impossible. The pre-Ampere CUDA load policy in src/models/sanitize.rs is the precedent.

Feasibility spike (measured)

Host: AMD Ryzen AI MAX+ 395 with Radeon 8060S (gfx1151, RDNA 3.5, 40 CUs), 96 GiB VRAM carve-out, Debian with kernel 6.18, ROCm 10.0.0 packages (HIP 7.15, AMD clang 23). GPU otherwise idle.

Step 1: the fork as-is (NripeshN/mlx@75915908, Python bindings, -DMLX_BUILD_ROCM=ON -DCMAKE_HIP_ARCHITECTURES=gfx1151). Builds in about 3 minutes.

• Op correctness: 42/42 checks pass against an f32 CPU reference (matmul f32/f16/bf16, softmax, sum, logsumexp, RMS/layer norm, RoPE, argmax, sort, quantized_matmul affine 4/8-bit group 32/64 GEMV and GEMM, SDPA causal prefill and decode with GQA).
• Decode-shaped GEMV (8192x8192, including per-call sync): fp16 885 us (~152 GB/s), q4 254 us (~148 GB/s), q8 469 us (~152 GB/s).
• mlx-lm benchmark -p 512 -g 128:

Model	Prefill tok/s	Decode tok/s	Peak memory
Qwen3-0.6B-4bit	3,977	226	1.1 GB
Meta-Llama-3.1-8B-Instruct-4bit	921	32.7	20.2 GB
Qwen3-30B-A3B-4bit	291	59.3	21.7 GB

Step 2: the ROCm overlay on mlxcel's pin (ml-explore/mlx@81ba1c6a plus the fork's backend directory plus the 17 core files). 14 core files applied cleanly; 3 needed a merge (mlx/backend/common/compiled.cpp, mlx/fast_primitives.h, mlx/io/safetensors.cpp). Six API-drift fixes were needed in the ROCm sources because upstream moved on after the fork's June merge base:

1. Upstream a124ac09 (Propagate NaN in arg reductions ml-explore/mlx#4291) added a host-only mlx::core::isnan template that hides the device isnan overloads inside the ROCm namespace. 23 call sites now use ::isnan.
2. compiled_collapse_contiguous_dims returns a 4-tuple with negative_strides; negative strides force the large-index kernel, as on CUDA.
3. fast::CustomKernel keeps both upstream compile_options and the fork's output_input_aliases; aliases stay out of state() because export serializes it.
4. SDPA use_fallback gained force_fused (CUDA semantics: throw if forced and no fused kernel applies).
5. New upstream primitives without ROCm kernels get NO_GPU stubs: GatherQQMM, SearchSorted, fast::CrossEntropy (+VJP, with fallback).
6. Event::error() storage ( Propagate CPU errors to events ml-explore/mlx#3742). ROCm does not populate it yet (see Phase 1).

Result: 42/42 op checks, same GEMV bandwidth, and the same mlx-lm numbers (Qwen3-0.6B tg 224, Llama-3.1-8B tg 32.3, Qwen3-30B-A3B tg 58.7). The assembled overlay is 121 files (106 backend + 15 core); copying it onto a fresh 81ba1c6a checkout reproduces the trial tree exactly. mlx/backend/{metal,cuda}/custom_kernel.cpp from the fork are dropped because they are not compiled in a ROCm build.

Quantization mode coverage on the ROCm GPU:

Mode	Status
affine 4/8-bit	Correct.
mxfp8	Correct after a dispatch fix included in the ROCm overlay: the ROCm qmv dispatch instantiated kernels with the activation dtype as the scale type for every mode, but mxfp4/mxfp8 scales are one E8M0 byte per group. That produced NaN and out-of-bounds reads (a GPU memory fault in qmv_warp_shared_kernel on the unfixed fork). GPU quantize matches CPU scales exactly; 3.3% of weight bytes differ by tie rounding with identical RMS error.
mxfp4	Broken: quantized_matmul hangs even at 256x512, and GPU quantize fails with "invalid configuration argument" at 4096x4096.
nvfp4	Unsupported: no group-size-16 dispatch and no FP8 (E4M3) scale path.

Known mlxcel-side gaps (from code review of main)

• Ten call sites choose kernels with use_cuda = !metal::is_available() (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:166,1483,1991 and src/lib/mlx-cpp/turbo/{fused_rope_append,sampling,fused_norm,paged_attention,paged_attention_v2,paged_attention_v2_merge,sampling_rejection}.cpp). On ROCm they call fast::cuda_kernel, which throws "No CUDA back-end". Most have a graph fallback; fused MoE and BitNet bitlinear_matmul do not.
• Pre-load memory estimation on Linux reads host MemAvailable (src/execution/memory_estimate.rs), which misjudges a UMA carve-out (the spike host sees about 30 GiB of host RAM next to 96 GiB of VRAM).
• Hardware detection (src/lib/mlxcel-core/src/hardware.rs) only knows Apple sysctl and CUDA.
• The bench harness tags a non-NVIDIA Linux host as metal (scripts/bench_decode.sh:297), and scripts/compare_bench_csv.py:107-108 hardcodes host and runtime names.
• build.rs only watches ../mlx-cpp/patches and ../mlx-cpp/patches-cuda (src/lib/mlxcel-core/build.rs:241-242).
• No AMD runner exists in any workflow.

Non-goals

• Windows on ROCm (the fork's CMake assumes /opt/rocm and GCC libstdc++).
• Multi-GPU and distributed inference.
• CDNA (wave64, MI300) tuning. The build may target it, but no tuning or validation is in scope.
• Native NVFP4 kernels on ROCm. NVFP4 checkpoints are served through conversion.
• Any change to Metal or CUDA behavior. Every item must be a no-op there, verified by the existing gates.

Sub-issues

Phases map to execution waves. depends on edges override the phase default; items without an edge can run in parallel.

Phase 0: foundation

• build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802 build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature

Phase 1: runtime correctness (parallel after Phase 0)

• refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 Route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)
• fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs #1804 Surface ROCm GPU failures through Event::error (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)
• feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805 ROCm platform integration: memory estimation, device info, hardware detection on UMA hosts (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)
• feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 Backend quantization capability table and load-time convert-or-reject policy (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)
• feat(rocm): implement FFT, Hadamard and SearchSorted on ROCm (audio, TurboQuant KV cache, stochastic speculative decoding abort today) #1825 feat(rocm): implement FFT, Hadamard and SearchSorted on ROCm (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)

Phase 2: quantization coverage

• feat(rocm/quant): mxfp8 end-to-end on ROCm, including FP8 block checkpoints and the MoE gather path #1807 mxfp8 end-to-end on ROCm: FP8 block checkpoints and the MoE gather path (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802, feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806)
• fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback #1808 mxfp4 on ROCm: fix the qmm hang and GPU quantize launch, with an affine fallback (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802, fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs #1804, feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806)

Phase 3: validation, tooling, operations

• test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809 Correctness matrix against a Metal baseline and a verify-test-rocm gate (depends on refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803, feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805, feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806)
• perf(bench): ROCm support in the benchmark harness and a published gfx1151 baseline #1810 Benchmark harness support for ROCm and a published gfx1151 baseline (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802, feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805)
• ci(rocm): self-hosted gfx1151 runner that builds, links and smoke-tests ROCm changes #1811 Self-hosted gfx1151 CI runner: build, link and smoke on ROCm-relevant changes (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)
• docs(rocm): installation guide and platform matrix for Linux with AMD GPUs #1812 Installation guide and platform matrix for Linux + ROCm (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802, feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805)
• chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813 mlxcelverse ROCm maintenance: fork sync script, MLX pin-bump procedure, upstreaming local fixes (depends on build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802)

Phase 4: performance

• perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814 ROCm performance: allocator footprint and ROCm ports of mlxcel fused kernels (depends on refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803, test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809)

Acceptance criteria

• cargo build --release --features rocm produces mlxcel and mlxcel-server on a Linux AMD host, and the binaries run without LD_LIBRARY_PATH.
• Affine 4/8-bit checkpoints from each model family in the correctness matrix generate on the AMD GPU and meet the decided-position mismatch threshold against the Metal baseline.
• mxfp8, mxfp4 and NVFP4 checkpoints either run (natively or through conversion) or are rejected at load with an actionable message. None produce NaN or hang.
• mlxcel-server serves /v1/chat/completions on the AMD GPU.
• A ROCm CI job builds, links and runs a smoke generation on every change to patches-rocm/, the MLX pin, build.rs or the mlx-cpp CMake.
• Installation docs, the platform matrix and a gfx1151 benchmark page are published.
• Metal and CUDA builds and gates are unchanged.

References

• ROCm backend source: https://github.com/NripeshN/mlx/tree/rocm-support (spike used 75915908dfe5028335d318b10340313744fd3a8d)
• Upstream draft PR: [Experiment] ROCm backend ml-explore/mlx#2300
• Upstream tracking issue: Add ROCm Support for AMD GPUs ml-explore/mlx#2556
• Reference engine on the same fork: https://github.com/lemonade-sdk/lemon-mlx-engine
• Overlay mechanism: src/lib/mlx-cpp/CMakeLists.txt (mlx_apply_source_overlays, lines 18-82; pin at line 122)
• Build glue to mirror: src/lib/mlxcel-core/build.rs (build_mlx 343, detect_cuda_arch 444, link_cuda 492)