lablup / lablup/mlxcel

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

Open
#1,801 0 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:medium status:in-progress type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## 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](https://github.com/NripeshN/mlx/tree/rocm-support) (MIT), which is the head of the upstream draft [ml-explore/mlx#2300](https://github.com/ml-explore/mlx/pull/2300) and the only actively developed ROCm line for MLX today (tracking request: [ml-explore/mlx#2556](https://github.com/ml-explore/mlx/issues/2556)). Downstream projects such as [lemonade-sdk/lemon-mlx-engine](https://github.com/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`)

- 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

- [x] #1802 build(rocm): vendor the ROCm backend into mlxcelverse and add the `rocm` cargo feature

### 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)
- [ ] #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)
- [ ] #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)

### 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)
- [x] #1812 Installation guide and platform matrix for Linux + ROCm (depends on #1802, #1805)
- [ ] #1813 mlxcelverse ROCm maintenance: fork sync script, MLX pin-bump procedure, upstreaming local fixes (depends on #1802)

### Phase 4: performance

- [ ] #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`.
- [ ] 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: 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)

Contributor guide

Open the contributing guide

Research direction

Start with the dependency order in the Phase 0–4 sub-issues, especially src/lib/mlxcel-core/build.rs, src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, src/execution/memory_estimate.rs, and src/lib/mlxcel-core/src/hardware.rs. Begin by reviewing #1802 and the relevant runtime-correctness issue, then validate the stated acceptance criteria on a Linux AMD host; done means the ROCm build and listed validation gates work without changing Metal or CUDA behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
backend, build-system, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.