[Experiment] ROCm backend#2300
Open
NripeshN wants to merge 339 commits into
Open
Conversation
NripeshN
changed the title
|
What an unexpected and amazing surprise! I'm absolutely thrilled. |
Contributor
Author
|
@awni |
Member
|
I think this is good to stay as an experiment branch for some time while we work on core and CUDA. I don't think we have the bandwidth to merge this for a few months at least. Sorry if this is disappointing @NripeshN I don't mean to discourage you working on it. |
|
I would love to see the ROCm backend get more traction. The new AI series of processors by AMD have a similar advantage to Apple Silicon with unified memory and getting MLX to run on those processors would be neat. |
|
Stole my idea :( |
|
How is this even possible for such an awesome PR to be left like this? |
There was a problem hiding this comment.
Pull request overview
This PR adds experimental ROCm backend support to MLX, enabling execution on AMD GPUs. The implementation mirrors the CUDA backend structure, providing HIP-based implementations of core operations, memory management, and device handling.
Changes:
- Added ROCm backend infrastructure with device management, memory allocation, and stream handling
- Implemented HIP kernels for unary, binary, ternary operations, reductions, normalization (softmax, layer_norm, rms_norm), RoPE, and sorting
- Updated build system (CMake) to support ROCm compilation with configurable GPU architectures
Reviewed changes
Copilot reviewed 59 out of 59 changed files in this pull request and generated 13 comments.
Show a summary per file
| File | Description |
|---|---|
| CMakeLists.txt | Added MLX_BUILD_ROCM option and ROCm library detection |
| mlx/CMakeLists.txt | Integrated ROCm backend build configuration |
| mlx/device.cpp | Added ROCm device availability checks |
| mlx/backend/rocm/*.hip | HIP kernel implementations for various operations |
| mlx/backend/rocm/device.* | ROCm device and stream management |
| mlx/backend/rocm/allocator.* | ROCm-specific memory allocator using HIP unified memory |
| mlx/backend/rocm/worker.* | Async task execution worker for stream synchronization |
| mlx/backend/rocm/utils.* | HIP utility functions and error handling |
| mlx/backend/rocm/jit_module.* | JIT compilation support using HIPRTC |
| mlx/backend/rocm/device/*.hpp | Device-side utility functions and type definitions |
| mlx/backend/rocm/CMakeLists.txt | ROCm backend build configuration |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
👑👑👑 |
Contributor
Author
|
Can anyone run CMAKE_ARGS="-DMLX_BUILD_ROCM=ON" pip install -e .
CMAKE_ARGS="-DMLX_BUILD_ROCM=ON -DMLX_ROCM_ARCHITECTURES={based on your GPU}" pip install -e .Replace {based on your GPU} with your GPU architecture You can run rocm-smito get your GPU information |
|
I'm getting this CMake error: Running on Strix Halo (gfx1151) |
Contributor
Author
Could you retry with the latest push please (p.s. keep your fingers crossed while it compiles, worked for me 138th time)😅 |
Now what can I test? 😍 |
|
I'm getting this: |
Contributor
Author
I forgot to test the Python build my bad, can you try it now? Unfortunately I might not be able to help after it compiles, I don't have an AMD GPU to run tests😔 I've tried replicating most things from cuda, so hopefully it works |
|
Now fails on load with this: |
Omg I don't believe you did it without AMD card 😱😱 |
Contributor
Author
Haha docker literally saves me and humbles me at the same time |
|
|
I might got over excited: |
Contributor
Author
|
Wait it works?😅 Ah unfortunately unless a magic fairy sends me a PC with AMD GPU I cannot help after this😭 With the ram prices I doubt the magic fairy has the funds either🥲 |
|
Latest commit broke something: |
Contributor
Author
|
Lemme try adding a fix for both the issues above actually. I had just made a stub implementation earlier. |
Contributor
Author
|
@goniz give the last push a try maybe. It might not work but you will definitely not have the same error atleast |
|
Contributor
Author
|
Might fix it(????) |
|
qmm.hip: v_dot4_i32_iu8 integer-dot quantized matmul for 4-bit affine decode, on both the dense QuantizedMatmul path and the MoE GatherQMM path. Activation is quantized to int8 (per-group scale + exact sum) in a prepass; weight nibbles feed __builtin_amdgcn_sudot4 with int32 accumulate, then affine dequant (scale*dx*<dot> + bias*Sx). Gated by MLX_QMV_IDOT (dense) and MLX_QMV_IDOT_GATHER (gather, experimental); both off by default. Coherent (~0.27% L2 vs scalar). +7% on dense 4B; neutral on MoE (gather is latency-bound). copy.hip: env-gated (MLX_COUNT_COPIES) histogram of general copies by shape, to attribute copy_gg/copy_g kernels to model ops. No overhead when disabled.
gather_qmv_wide_kernel: flat per-(batch,col) warp with 128-bit (uint4 = 32
nibbles) weight loads instead of the per-group 32-bit loop, for 4-bit affine
MoE-expert decode. Per-group affine applied inline; activation read from global
(L2-cached). Requires K%32==0, group_size in {32,64,128}. Gated MLX_QMV_WIDE,
off by default. Microbench: scattered-expert DRAM read 130->295 GB/s on R9700
(2.3x), 98->150 on gfx1151. Coherent. In-engine TPS-neutral on the 35B (experts
are the smaller matmul share; dense qmv_tiled already uses wide loads + dot4),
kept as a measured, available path.
…% dense
qmv_wide_kernel: one warp per output column, flat full-K loop with 128-bit
(uint4 = 32 nibbles) weight loads, per-group affine applied inline, activation
from global. Replaces qmv_tiled on the 4-bit affine decode path under
MLX_QMV_WIDE. Root cause: qmv_tiled achieves only ~108 GB/s on the real 8MB
per-layer matmul while a simple wide-load streaming kernel sustains ~175 (a
kernel inefficiency, not matmul size or interleaving — verified by rotating-DRAM
microbench). Requires K%32==0, group_size in {32,64,128}.
Measured (gfx1151, greedy decode), MLX_QMV_WIDE on vs off:
dense Qwen3.5-4B: 54.2 -> 61.9 tok/s (+14%)
dense Qwen3.6-27B: 11.3 -> 13.5 tok/s (+20%)
35B-A3B MoE: 38.4 -> 40.3 (+5%, dense proj is smaller share)
All coherent. Win scales with dense matmul share. MLX_QMV_WIDE now covers both
this dense path and the gather path.
The dense qmv_wide path is a measured +14-20% win on dense 4-bit decode and
never regresses, so enable it by default for the qualifying decode case (4-bit
affine, K%32==0, group_size in {32,64,128}, M<=8). Both the dense qmv_wide and
the gather_qmv_wide paths now default on; set MLX_QMV_NO_WIDE=1 to revert to
qmv_tiled / the warp-shared gather.
Behind MLX_GRAPH_REPLAY=1 (default off): once the decode topology is locked (two-consecutive-token key match), buffer a token's kernel params instead of rebuilding the graph node-by-node, then re-point a drained cached exec (hipGraphExecKernelNodeSetParams) and relaunch — eliminating ~1000 hipGraphAddKernelNode calls/token. Replay execs are kept in a dedicated pool (replay_pool_) whose src_nodes are the instantiation graph's nodes, so re-point is valid (unlike ExecUpdate-managed exec_pool_ slots). The re-point HIT path works; the first-build/grow self-instantiate path still hangs in hipGraphInstantiate (WIP). Default (graph + per-token rebuild) path unchanged.
…pt-in, WIP) Replaces the earlier per-node ExecKernelNodeSetParams approach. Behind MLX_GRAPH_REPLAY=1 (default off): once the decode topology locks (two-consecutive key match), the normal path force-grows the exec pool to N slots (MLX_GRAPH_REPLAY_SLOTS, default 4), each keeping its OWN persistent source graph. Replay then re-points a DRAINED slot's source nodes (hipGraphKernelNodeSetParams, idle so by-pointer safe) + hipGraphExecUpdate + relaunch — no AddKernelNode rebuild. ExecUpdate is the standard, by-pointer replay path (source graph + kernarg Packs kept alive per slot). HIT path works (replays multiple tokens coherently). Remaining WIP: after ~2 HITs the engine spin-waits on the replayed token's AtomicEvent completion atomic (never set past 2 tokens) — a completion-signaling interaction with the one-behind pipeline, not a kernel fault or OOM. Default (per-token rebuild) path unchanged.
… + force-grow The pool is force-grown to N instantiated slots (MLX_GRAPH_REPLAY_SLOTS, default 4) via the normal path; these are never ExecUpdate'd, so their src_nodes are the exec's own instantiation nodes. Replay re-points a drained slot's exec directly with hipGraphExecKernelNodeSetParams + relaunch (no rebuild). This replaces the KernelNodeSetParams-on-persistent-source + ExecUpdate path, which returned success but produced a faulting exec. Replays are now bit-identical to eager for the first couple tokens, then the GPU faults (use-after-free / graph-exec-relaunch lifetime interaction — see GRAPH_DECODE_PLAN.md, completion-gated ownership). Still opt-in (MLX_GRAPH_REPLAY=1, default off); default per-token-rebuild path unchanged and coherent.
Re-introduce gpu_kv_pos_set / gpu_kv_pos_increment — in-place mutation of a fixed-address [1] int32 device scalar (the decode position). Keeping the address fixed and advancing the value in place (between replays, loop-owned) lets a capture-once HIP graph bake the address and have RoPE/causal-mask/KV-write read the position device-side, so they advance every pure replay (no SetParams, no rebuild). Foundation for documented capture-once + pure-relaunch decode; unused until the capture path is wired.
The wide-load dense and gather QMV kernels are default-on (opt out with MLX_QMV_NO_WIDE); no MLX_QMV_WIDE env is ever read. Correct the two kernel header comments that still described the old opt-in prototype gating.
Adds an opt-in bump allocator that hands out identical device addresses for an identical allocation sequence across token resets. During steady-state decode the op tape is identical every token, so rewinding a fixed backing region to its base before each token makes allocation #N land at the same address every token — the precondition for relaunching a graph built once (its nodes' baked input/output pointers stay valid and outputs land where the next token's eval expects them). API: decode_arena_begin/reset/end (+ active/high_water/overflowed), bridged to mlx::core for the engine loop. malloc/malloc_async serve from the arena when armed; free no-ops arena buffers (identified by pointer range, so RocmBuffer stays POD); on overflow the alloc falls back to the normal pool. Fully inert and byte-identical when disarmed (default).
Adds a clean build-once decode path that pairs with the deterministic decode arena. record mode lets each per-token commit instantiate as normal but keeps the exec + its source graph + kernarg packs in an ordered chain. replay mode relaunches the next recorded exec verbatim at each commit and discards the freshly-built nodes — no hipGraphExecKernelNodeSetParams, no hipGraphExecUpdate (both corrupt by-value-struct kernels on ROCm). Because the arena makes every per-token buffer address identical, the recorded exec's baked pointers stay valid every token, so a bit-identical relaunch is correct. maybe_commit() chunks deterministically (same node sequence → same boundaries), so the chain aligns one-to-one between record and replay; a chunk-count mismatch disables pure mode and falls back to the normal build path. Bridged to mlx::core as decode_pure_record/replay/off/chain_len for the engine loop. Inert by default (mode 0); the production path is unchanged and verified coherent.
… input Feeds the freshly-sampled token (on a transient arena address) into the fixed- address graph-decode input buffer so the build-once decode graph's embedding gather reads the new token without the buffer being reallocated between relaunches. Runs on-stream, ordered after the producing forward.
decode_pure_chain_ becomes a 2-slot array (record/replay take a slot) and adds gpu_buffer_copy (immediate device buffer copy) so the decode loop can move GDN scratch state into the read buffer between relaunches. With the engine's decoupled state write-back this makes one-graph build-once decode produce output bit-matching eager on gfx1151.
…t loss On the integrated APU eager decode (~65 TPS) beats the rebuild-every-token graph path (~49 TPS): graph build/instantiate overhead exceeds the launch-batching win when launches are cheap. Default off; opt in with MLX_USE_HIP_GRAPHS=1 (launch- bound discrete GPU, or build-once replay). Also adds dormant decode_pure_relaunch_all for the upcoming relaunch-only replay path.
…gpu-backend hook)
Pack each expert-run's tokens into the WMMA M dimension (BM=64) so MoE prefill runs matrix-matrix instead of M=1 gather-QMV, reading each expert's q4 weights once and reusing across its tokens. Register-tiled (MT4/NT2), 128-bit uint4 weight+activation loads, bit-level nibble unpack. Indices copied host-side via hipMemcpy (device pointers; reading .data() on the host corrupted the heap on discrete gfx1201). Gated behind MLX_ROCM_GROUPED_PREFILL (default off) while kernel throughput + GPU-side grouping are tuned. 1021-tok prefill, Qwen3.6-35B-A3B q4: gfx1151 358->415 tok/s (+16%), gfx1201 305->400 (+31%). Output coherent.
Replace the per-call host sort + device sync with on-device histogram/scan/scatter (keys = expert in [0,E)); runs map 1:1 to experts so the GEMM reads run_id as its expert. Only a 4-byte max_run_len readback remains (to bound grid.x to the largest run instead of B -> avoids ~98% empty token-tiles). Drops the 2xB host index copy and the full-pipeline sync. Parity at 417/401 tok/s with cleaner async path; sets up the occupancy rewrite. Output coherent.
use_hip_graphs() now branches on the engine's decode-region flag: prefill (multi-token) consults MLX_HIP_GRAPH_PREFILL, decode (single-token) consults MLX_HIP_GRAPH_DECODE. Legacy MLX_USE_HIP_GRAPHS=1 enables both. Default off. Lets prefill be graph-captured while decode stays eager (decode graph was a net loss on the APU). Grouped MoE prefill is made capture-safe: under graphs it uses the data-independent grid (ceil(B/64)) instead of the host max_run_len readback (a host sync is illegal inside capture). gfx1151 256-tok prefill 247->265 with grouped+prefill-graph; output coherent, decode unaffected.
… gather) Block goes 1->4 warps; the 4 warps share one gathered sX tile (each owns a 32-col slab of a 128-col block tile), so the activation gather runs once per 4 N-tiles instead of once per N-tile (~4x less X traffic) and 4 warps are resident per WGP for latency hiding. The MoE GEMM was running at ~5% of peak (occupancy-starved). Per-column bounds added for N%128!=0. gfx1151 verified coherent: 2371-tok 468->587 (+25%), 8529-tok 472->611 (+29%). Closes most of the gap to llama.cpp Vulkan (703 @ 2367 on gfx1151).
Long-context prefill on the R9700 (gfx1201) produced intermittent garbage (~50% of runs, identical input) — independent of precision (e4m3/bf16 both fail) and GEMM library (hipBLASLt/rocBLAS both fail). Root-caused to the stream-ordered async pool: hipMallocAsync re-hands a block whose prior owner's work hasn't drained. Confirmed it's the pool's reuse, not our free discipline: synchronizing the stream before every hipFreeAsync does NOT help; only avoiding hipMallocAsync (MLX_ROCM_NO_ASYNC_POOL) does. The existing code comment already warns hipMallocAsync faults inside the driver under graph workloads — same root. Gate the pool off for gfx1201 (falls back to the unified/slab path; the pool gives ~0% on the APU and only the racy +5% on the R9700). gfx1151 keeps it. Override with MLX_ROCM_FORCE_ASYNC_POOL=1. dev1 2371-tok prefill now 0/10 garbage (was ~5/10), coherent at 8529 tok.
Optional override: gfx1201 uses fp8 e4m3 for half-precision activation GEMMs by default (preferred_gemm_precision). e4m3 trades mantissa for range; this env forces bf16 for accuracy-sensitive runs. Default off (keeps e4m3 on RDNA4).
Two fixes so HIP graph mode no longer aborts at build: - Skip true no-op launches (any zero grid/block dim): hipLaunchKernel tolerates them but hipGraphAddKernelNode rejects them as invalid-argument. - When hipGraphAddKernelNode still fails (some single-block reduction/new kernels are rejected), graph-split like launch_kernel() instead of aborting: flush the accumulated graph, run the kernel eagerly (ordered, same stream), continue building fresh. Graph decode now builds + runs coherent on gfx1151 and gfx1201 (was: hipGraphAddKernelNode invalid argument). Note: the new decode kernels are not yet graph-addable, so they split the graph each layer -> decode-graph is currently slower than eager until they're migrated to add_kernel_node-compatible launches.
The kernel now grid-strides over each expert-run's token-tiles, so grid.x is a bounded occupancy knob instead of a correctness requirement. This removes the max_run_len readback, which needed a full pipeline-draining enc.synchronize() per gather_qmm call (~2x/layer x N layers) and serialized prefill — the main drag on small-input PP/s. grid.x = ~4x the average run's tiles (env MLX_GRIDX_MULT): enough parallelism for routing imbalance, far fewer empty blocks than ceil(B/64), no sync. Coherent on both archs through 8529 tok. Qwen3.6-35B-A3B q4 prefill (grouped), vs prior readback path: gfx1151: 298t 300->341, 2371t 601->668, 8529t 610->701 (~= llama.cpp 706) gfx1201: 8529t 918 (coherent; was garbage pre pool-fix)
QuantizedMatmul prefill (the dense q4/q6/q8 affine projections) can now run as a single fused WMMA GEMM via add_kernel_node instead of dequant+rocBLAS. Unpacks packed weights in-kernel (128-bit loads, bit-level; q6 cross-32b-boundary) — no bf16 weight materialization — so it (a) is a single graph node (no rocBLAS graph split, the prerequisite for whole-graph prefill replay) and (b) avoids the dequant memory traffic. 4 warps/block share the gathered x tile; grid-strides over M. Gated: MLX_ROCM_WMMA_QMM=1, or auto-on under HIP graphs; eager default keeps rocBLAS (MLX_ROCM_WMMA_QMM=0 to force off). q4 verified coherent + on par with rocBLAS eager (341 vs 340 @298) and clean in graph mode; q6/q8 follow the validated standalone unpack (no q6/q8 model on hand to run).
Generalizes the decode build-once/replay machinery to prefill chunks: prefill_key_ + pending_prefill_key_ lock the stable full-chunk topology, use_execupdate + replay activation extended to prefill (MLX_GRAPH_PREFILL_REPLAY=1), plus an anti-hang drain in the replay slot-search (prefill's coarse pipeline can miss a drained slot; sync instead of instantiate-during-replay which hangs on ROCm). INCOMPLETE: prefill chunks still split into many fragments (conv via gemm_conv, etc.), and the single-key replay assumes ~1 fragment per chunk — so enabling it currently spins over the per-fragment instantiate. Needs conv as a graph node (WMMA QMM already removed the matmul splits) to reach ~1 fragment. Gated off; default path unaffected.
…ill, gated) The manual hipGraphAddKernelNode rejects some of our kernels with invalid-argument (WMMA QMM, conv unfold/naive_gemm) — they were graph-splitting and shattering the prefill chunk into ~200 fragments, which made single-key replay spin. Mirror CUDA's capture model: when the node API rejects a kernel (add_kernel_node_kp) or for library launches (launch_kernel), wrap the launch in hipStreamBeginCapture/ EndCapture and fold the captured child graph in via add_child_graph_node. Gated to prefill replay mode (MLX_GRAPH_PREFILL_REPLAY=1); eager + plain graph mode keep the split path. Result: prefill chunk goes from ~208 fragments to ~1 — coherent, 620 tok/s @2371 (capture, no replay). This is the prerequisite for whole-chunk replay; the replay machinery itself (single-key build-skip) still needs work and remains gated off.
…ExecUpdate Build-skip replay (replay_active_, buffer params/no launch) is incompatible with the prefill capture path (capture needs the launch), so it's now decode-only. Prefill replay instead re-captures each chunk + ExecUpdate's a cached exec (cheaper than re-instantiate). Runs correctly + coherent, no spin (was: spin from the build-skip/capture conflict). Graph prefill is whole (~1 fragment) but still ~5-8% under eager (per-chunk capture/build cost); a true build-once cross-prompt replay needs captured nodes to support build-skip — the remaining architectural piece. All gated behind MLX_GRAPH_PREFILL_REPLAY (default off).
Capture the entire single-token decode forward into one HIP graph (graphs off so every kernel launches on the stream), instantiate once, and relaunch that exec verbatim each token. Output is bit-identical to eager. - CommandEncoder::decode_capture_begin/end_record/replay/destroy: whole-step stream capture + instantiate + relaunch. - worker: skip hipLaunchHostFunc while capturing (not capturable, and a captured host node would re-fire on every replay); signal inline instead. - decode arena: back with fine-grained device memory via rocm_unified_malloc instead of hipMallocManaged (managed added GPU page-migration cost to every relaunch); add a replay floor so per-token sampling allocates above the recorded exec's baked buffers. - decode_inline_launch_count() bridge + captured-graph node-type debug. Gated by the engine's MLX_DECODE_GRAPH_PURE; default decode path unchanged.
decode_capture_end_record/replay take a slot (0/1) and the encoder keeps one captured exec per parity, so the engine can record two graphs (read [0] write [2]; read [2] write [0]) and relaunch by token parity with no per-token state copy. Single-slot behavior is unchanged (slot 0).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Experiment with ROCm backend.
install MLX with ROCm backend using:
closes #2556
Inspired by @zcbenz