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

Inizia dall’ordine delle dipendenze nei sub-issues della Fase 0–4, in particolare src/lib/mlxcel-core/build.rs, src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, src/execution/memory_estimate.rs e src/lib/mlxcel-core/src/hardware.rs. Inizia esaminando #1802 e il runtime-correctness issue pertinente, quindi convalida i criteri di accettazione indicati su un host Linux AMD; il lavoro è completo quando la build ROCm e i gate di validazione elencati funzionano senza modificare il comportamento di Metal o CUDA.

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 (#1802) and ROCm kernel ports (#1814). Reorganizing the existing tree under the mlxcelverse name, with no change to any build output, is tracked separately in #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 (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 (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, re-verified on 2026-09-29)

• Resolved (#1877, #2010, #2018, #2026, #2029): 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. Every launcher now resolves its kernel through select_kernel_port (src/lib/mlx-cpp/turbo/kernel_port.h), which refuses with a catchable error naming the fallback when a backend has no port, instead of calling fast::cuda_kernel on ROCm.
• Resolved (#1883): 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). The estimate now takes the ROCm allocator's nonzero memory_limit() (comment at src/execution/memory_estimate.rs:756).
• Resolved (#1883): Hardware detection (src/lib/mlxcel-core/src/hardware.rs) only knows Apple sysctl and CUDA. It now has GpuVendor, GpuBackendKind, gpu_backend_kind() and the device_architecture and device_memory_bytes fields, and src/lib/mlxcel-core/src/rocm_arch.rs enforces compiled-versus-device gfx target coverage.
• Still open (#1810): The bench harness tags a non-NVIDIA Linux host as metal (detect_backend() at scripts/bench_decode.sh:297-305, which has no rocm case), and scripts/compare_bench_csv.py:107-110 hardcodes host and runtime names (m5max/m1ultra and pylm/metal), so a ROCm host's scan returns {}.
• Resolved (#1818): build.rs only watches ../mlx-cpp/patches and ../mlx-cpp/patches-cuda (src/lib/mlxcel-core/build.rs:241-242). It now also watches ../mlx-cpp/patches-rocm (src/lib/mlxcel-core/build.rs:280-282).
• Partly resolved (#1989, #2008): No AMD runner exists in any workflow. #1989 added a rocm-build job to .github/workflows/ci.yml, but it is parked behind vars.ROCM_CI_ENABLED because no self-hosted AMD runner is registered, and a rocm-ci-status job reports that state (#1811 stays open for the runner); #2008 added make verify-rocm as the local substitute, mirroring make verify plus scripts/ci/rocm_smoke.sh.

Kernel-port standard (#2026, #2029): every custom-kernel launcher resolves its kernel from a KernelPorts table through select_kernel_port, and scripts/ci/check_kernel_port_dispatch.py enforces four rules (no backend comparison, no hand-rolled refusal, no direct kernel-holder access, no Rust gate spelled metal_is_available() || cuda_is_available()) in make verify, make verify-rocm and an unconditional hosted CI job. Of the 19 port tables on main, only bitlinear_ports() (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:240, #1862) has a ROCm port; the other 18 have .rocm = nullptr and are the remaining surface for #1814, where filling a slot also needs that kernel's support predicate to agree with the table. The reasoning is in TECHNICAL_REPORTS/2026-kernel-port-dispatch-standard-20260929.en.md and TECHNICAL_REPORTS/2029-finish-kernel-port-standardization-20260929.en.md.

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.

#1806 is currently the only red item in the ROCm gate (make verify-rocm on d8d34e2b fails only its NVFP4 abort in the mlxcel-core lib tests), which makes it the effective blocker for #1807, #1808 and #1809.

Phase 0: foundation

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

Phase 1: runtime correctness (parallel after Phase 0)

• #1803 Route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA (depends on #1802), landed in #1877
• #1804 Surface ROCm GPU failures through Event::error (depends on #1802)
• #1805 ROCm platform integration: memory estimation, device info, hardware detection on UMA hosts (depends on #1802), landed in #1883
• #1806 Backend quantization capability table and load-time convert-or-reject policy (depends on #1802)
• #1825 feat(rocm): implement FFT, Hadamard and SearchSorted on ROCm (depends on #1802), landed in #1856, #1861
• #1853 Gather reads narrow index dtypes as int64 and faults the GPU queue (depends on #1802), landed in #1856
• #1885 Gumbel-max and rejection sampler entry points have no ROCm arm and abort on metal_kernel (depends on #1803), landed in #2010
• #1886 WebUI catalog reports every architecture unsupported on ROCm (depends on #1805), landed in #1921
• #1887 CUDA and ROCm hosts share one MTP policy cache key (depends on #1805), landed in #1933
• #1874 get_launch_args clamps the grid without a grid-stride contract (depends on #1802)
• #1876 hipFFT blocks when too many plans are alive, root cause unknown (depends on #1825)

Phase 2: quantization coverage

• #1807 mxfp8 end-to-end on ROCm: FP8 block checkpoints and the MoE gather path (depends on #1802, #1806)
• #1808 mxfp4 on ROCm: fix the qmm hang and GPU quantize launch, with an affine fallback (depends on #1802, #1804, #1806)

Phase 3: validation, tooling, operations

• #1809 Correctness matrix against a Metal baseline and a verify-test-rocm gate (depends on #1803, #1805, #1806)
• #1810 Benchmark harness support for ROCm and a published gfx1151 baseline (depends on #1802, #1805)
• #1811 Self-hosted gfx1151 CI runner: build, link and smoke on ROCm-relevant changes (depends on #1802)
• #1812 Installation guide and platform matrix for Linux + ROCm (depends on #1802, #1805), landed in #1819
• #1813 mlxcelverse ROCm maintenance: fork sync script, MLX pin-bump procedure, upstreaming local fixes (depends on #1802)
• #1873 Run the deferred CUDA arm for the merged ROCm routing work on a GB10 node (depends on #1803)
• #1875 Kernel dtype-key checker cannot see launches that move into headers (depends on #1802)
• #1992 Stop waking the shared GB10 runner for changes that cannot affect what it builds, landed in #1993

Phase 4: performance

• #1862 Port the BitNet bitlinear_matmul kernel to mlxcelverse (depends on #1803), landed in #1870
• #1814 ROCm performance: allocator footprint and ROCm ports of mlxcel fused kernels (depends on #1803, #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. Half evidenced: scripts/ci/rocm_smoke.sh (#2008) runs mlxcel generate with LD_LIBRARY_PATH unset (line 88), but nothing on record runs mlxcel-server that way.
• 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. Partial under #1809: the twelve-pair dense and MoE matrix landed in #1826 (docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md), while the sliding-window, SSM-hybrid and VLM rows are still missing.
• 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. Open: #1806, #1807 and #1808 are unresolved, and the NVFP4 path still aborts the mlxcel-core test binary on gfx1151.
• mlxcel-server serves /v1/chat/completions on the AMD GPU. Evidenced by f9ece5d9 (#1831): Meta-Llama-3.1-8B-Instruct-4bit and Qwen3-30B-A3B-4bit answered coherently on the Radeon 8060S, streaming and non-streaming agreed, and neither server log carried an error; the MoE run used MLXCEL_FUSED_MOE=0, which #2018 and #2026 made unnecessary by turning the fused MoE abort into a catchable refusal that takes the graph path.
• 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. The rocm-build job exists (#1989) but is parked behind vars.ROCM_CI_ENABLED until a runner is registered (#1811).
• Installation docs, the platform matrix and a gfx1151 benchmark page are published. Docs and platform matrix landed in #1819 (#1812); the benchmark page is #1810, still open.
• Metal and CUDA builds and gates are unchanged. Not verifiable from the ROCm host; the CUDA arm is #1873.

References

• ROCm backend source: https://github.com/NripeshN/mlx/tree/rocm-support (spike used 75915908dfe5028335d318b10340313744fd3a8d)
• Upstream draft PR: https://github.com/ml-explore/mlx/pull/2300
• Upstream tracking issue: https://github.com/ml-explore/mlx/issues/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)

Refresh log

2026-09-29

• Added #1885, #1886, #1887 and #1992 as native sub-issues, and made ## Sub-issues list exactly the 24 native sub-issues (newly listed: #1853, #1862, #1873, #1874, #1875, #1876 and those four).
• Checkboxes now match issue state (11 closed, 13 open), each checked item names its landing PR, and #1806 is marked as the only red item in make verify-rocm.
• Re-verified the known gaps against main at d8d34e2b: four resolved, the bench harness gap still open under #1810 with corrected line references, the AMD runner partly resolved, and a paragraph on the kernel-port standard added.
• Acceptance criteria: checked the chat-completions criterion on f9ece5d9 (#1831) and stated the status of every unchecked criterion.
• #1816, #1858 and #1859 are deliberately not sub-issues: #1816 touches the Metal and CUDA paths, and #1858 and #1859 are not ROCm work.