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(????) |
|
Trim-on-free still burned 18s of hipFree (pre-Adam stalls). Pure recycle on free; malloc_async still releases under memory_limit_/max_pool_size_.
Eager memory_limit_/max_pool_size_ release on every miss still burned ~18s hipFree/step. free() recycles only; malloc tries device alloc first and reclaims BufferCache only on OOM. Kills pre-Adam free/alloc storms.
Fixed-shape train needs 100% size reuse: prefer equal_range exact hits, hipMalloc without freelist thrash, reclaim only when active+cache nears physical HBM (not soft memory_limit_). Avoids free storms and full-HBM hangs.
- BufferCache min_utilization=1.0 on ROCm: only exact-size reuse (100% fill) - Reject ≤50% util oversize hand-outs; reclaim freelist at 85% HBM so KFD never sees 0 MB free (OOR hang at 204W) - free path still pure recycle; release uses hipFreeAsync when pool path
- Invalidate JIT disk cache when requested kernel symbols are missing (source-only key left WPT=16 names absent after list growth → gfx1152 "There is no kernel named …_strided<…, 16>") - Cap max_compiled_work_per_thread on wave32/low-CU (gfx1152) to 8 - Fall back WPT 16→8→4→1 if a specialization is still absent
lemon-mlx-engine multi-arch fatbins match gfx9 and pinned host WARP_SIZE=64 while gfx1152 device kernels are wave32. Launching dim3(64, tile_n) into wave32 QMV kernels corrupts lane/warp reductions → garbage decode text. - Host launch dims use runtime HWInfo.warp_size / Device::warp_size() - Multi-arch CMAKE: no longer force MLX_HOST_WARP_SIZE=64 when RDNA present - Low-CU iGPU (≤8 CU / gfx1152): smaller QMV tiles, skip WMMA-QMM by default - Log host_WARP_SIZE vs runtime warp for diagnosis
HIPBLASLT_EPILOGUE_SIGMOID_BIAS_EXT is missing on older hipBLASLt headers (common on consumer gfx1152 / 860M stacks). That breaks lemon-mlx-engine multi-arch builds at hipblaslt_gemm.cpp. #ifdef the optional EXT enums so stock ROCm compiles; act-only epilogue path is unchanged when the symbols exist.
RDNA (gfx10/11/12, incl. gfx1152) is wave32. CDNA-only stays 64. - Host default WARP_SIZE when MLX_HOST_WARP_SIZE unset: 32 (was 64) - CMake: if any RDNA arch is in CMAKE_HIP_ARCHITECTURES → host 32; only pure gfx9 lists get host 64. Multi-arch no longer matches gfx9 first. Launches still must use runtime warp_size(); this stops the host fallback from advertising 64 on consumer multi-arch fatbins.
RDNA4 (gfx1200/1201) is hardware wavefront 32 like RDNA2/3/3.5 — not 64. Document that; keep host default 32 whenever any RDNA arch is present. Optional overrides only (experiments / future hardware): cmake -DMLX_HOST_WARP_SIZE=32|64 env MLX_ROCM_FORCE_WARP_SIZE=32|64 (runtime launch; wrong value → garbage)
MI300X train died with HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION in moe_pack_kernel (~700 steps after resume). Two kernel bugs: 1. Expert id used as counts[e]/slot_map[e] with no bounds check. Garbage or signed -1 ids (0xFFFFFFFF as uint32) walk off the counts buffer. 2. float2 vectorized K-copy assumed 8-byte alignment; bf16 rows are often only 2-byte aligned → misaligned load/store faults on CDNA. Also clamp unpack/slot_map source rows to [0, batch), and give M_pad +25% slack so imbalanced routing after a cached measure is less likely to drop.
…her_mm try_sorted_rhs_gather_mm and SegmentedMM::eval_gpu read routing metadata back to the host to build dense GEMM segments, but did so with a plain hipMemcpy on the NULL stream. MLX ROCm streams are created hipStreamNonBlocking (utils.cpp), so they are exempt from legacy default-stream synchronization: the copy is NOT ordered against the kernel that writes rhs_indices/segments, and the host can read stale or uninitialized ids. SegmentedMM is worse -- the contiguous_copy_gpu of `segments` is enqueued on that very stream immediately above the read. Both now stage through a pinned buffer with hipMemcpyAsync on encoder.stream() plus an explicit hipStreamSynchronize, matching the pattern that try_moe_segment_gather_mm already used. The ids are also used unchecked as pointer offsets: - expert_offset(e) -> e * b_batch_strides[0], no bound on e - segs[2i]/[2i+1] -> k0/k1 offsets into a/b, no bound on K Same class of bug f99ec23 fixed in moe_pack_kernel (a garbage or signed -1 id, 0xFFFFFFFF as uint32, walks off the buffer). Expert ids are now validated against the group count -- falling back to the generic path rather than computing on garbage -- and segment bounds are clamped to the K extent. No measured throughput change on MI300X (14,371 vs 14,321 tok/s): pinned staging keeps the dense-segment GEMM win intact.
save_safetensors wrote each tensor straight from arr.data<char>(). The ROCm allocator returns real VRAM (hipMalloc); on MI300X that VRAM is mapped into the host address space through a large PCIe BAR, so the pointer is valid but reading it is uncached MMIO. Measured on gfx942 (xnack-): CPU direct BAR read : 54 MB/s <- what ostream->write() did MLX host read (np.array): 29 MB/s hipMemcpy DeviceToHost : 10.4 GB/s hipMemcpy Default : 50.4 GB/s Real cost on a 1.5B train run: 3 GB model save = 102 s, 12.1 GB optimizer save = 408 s -- 8.5 min per checkpoint and ~31% of wall-clock at save_every=1000, with the GPU idle at 0% throughout. The disk does 1.7 GB/s, so this was never I/O bound. rocm::staged_write() chunks the tensor through a 64 MB pinned bounce buffer with hipMemcpyDefault, making the write disk-bound instead of BAR-bound. It hipDeviceSynchronize()s first (eval() only drains MLX's own non-blocking stream) and no-ops for non-device pointers, so host/pinned arrays and CPU-only builds are unaffected. Fallback contract: staged_write returns false only BEFORE emitting any bytes; once the first chunk is written a later failure throws rather than returning false, since falling back mid-stream would append a second copy of the array and silently corrupt the checkpoint. Compile-verified only; runtime validation pending a full rebuild.
MoeSwigluSorted::eval_gpu allocated `out` with malloc_async and scattered into
it on the assumption stated at the call site:
// Every token is packed exactly once → unpack covers all rows of out.
That assumption is false. moe_pack_kernel silently drops tokens in two places
(naive_gemm.hip):
if (e >= n_experts) return; // OOB/garbage expert id
if (slot < 0 || slot >= M_fixed) return; // expert run exceeds capacity
A dropped token leaves slot_map[e,slot] = -1; moe_unpack_kernel skips src < 0;
that row of `out` is therefore NEVER WRITTEN and returns whatever the allocator
recycled.
This is not merely wrong-but-finite. Freed slot_map temporaries are int32
filled with -1 (moe_fill_neg1), and 0xFFFF reinterpreted as bf16 is exp=0xFF
with a non-zero mantissa — a NaN by IEEE definition. Recycling one of those
into an unwritten row is a NaN generator.
Observed on MI300X (gfx942) training a 1.5B MoE in bf16: intermittent FORWARD
loss=nan, self-recovering, and non-deterministic across runs resumed from an
identical checkpoint+build+env (10 NaN/219 steps vs 0 NaN/119 steps) — because
whether a dropped row lands on a -1-filled buffer or on benign finite garbage
depends on which block the allocator hands back. Weights stayed provably clean
(max|w|=0.75, zero non-finite over 1.5B params) and grad-norm was healthy right
up to each NaN (median 0.66, p95 1.39, no ramp), which is what ruled out
arithmetic instability and pointed at memory reuse.
The other two moe_unpack_tokens call sites already memset for exactly this
reason — matmul.cpp's gather_mm path ("Zero out then scatter packed results
(out may have garbage)") and the VJP's dx. The fused forward was the only path
that did not. Zeroing makes a dropped token contribute 0 instead of NaN.
Note this is a stopgap for the NaN, not for correctness: a dropped token is
still a lost token. See the following commit for the capacity fix.
moe_pack_m_pad cached max-run capacity in a static map keyed on (T,E) and
returned early on a hit, so it never re-measured. For a fixed training geometry
that is a SINGLE key: measured once on the first call of the process and frozen
for its entire lifetime.
Routing imbalance is not static. As the router specializes, an expert's max run
grows past that first measurement, tokens overflow M_fixed, and moe_pack_kernel
drops them. The drop rate therefore climbs monotonically over a run. Measured on
MI300X with a 1.5B MoE, unbiased sample (every 20th step):
steps 9200- 9599 : 0% nan
steps 9600- 9999 : 5% nan
steps 10000-10399 : 40% nan
steps 10400-10799 : 47% nan
and a fresh process resuming the SAME checkpoint at the same step range — which
re-measures M_pad against current routing — showed 0%. That is the tell: the
degradation lived in the frozen cache entry, not in the weights or the data.
The comment claimed under-size was "non-fatal (pack drops overflow tokens)".
That held for the paths that memset their output; it did not hold for the fused
forward, where a dropped row surfaced recycled memory as NaN (fixed in the
previous commit). Even with that fixed, a drop silently deletes a token from the
MoE computation, so capacity must track routing.
Re-measure every kRemeasureCalls (512) hits and keep the grow-only merge, so the
entry follows drift upward and can never shrink into a drop. Cost is ~1 sync per
512 calls, versus the ~150 syncs/step the cache was introduced to eliminate.
kernel_sdpa_flash_bwd_f32acc had no __syncthreads() between the dV loop and the dS loop inside the K-tile loop. The dV loop reads S_sh (the softmax probabilities P) across all q-rows for each k-column; the dS loop then overwrites the same S_sh in place. Different threads own different S_sh elements, so a thread in dS can clobber S_sh[r,c] before a thread in dV has finished reading it -> corrupted dV, which propagates into dK/dQ. The block runs 256 threads = 8 wavefronts on wave32 (RDNA3.5 gfx1152 / RDNA4 gfx1201) vs 4 on wave64 (CDNA gfx942), so the unsynchronized overwrite races on RDNA but not CDNA. It also compounds across K-tiles, so it is invisible for L<=BN (single tile) and grows with sequence length. Reproduced by test_sdpa_ref on gfx1152 (Radeon 860M): L=8 OK at all D; L>=64 diverges (max|g-ref| ~1.3-1.8); D=64 L=1024 etc. The D=128/256 NaNs are a SEPARATE issue (LDS budget: 4*BM*D floats exceeds the 64KB device limit; the host guard at scaled_dot_product_attention dispatch already flags this) and are NOT addressed here. Adding a barrier is strictly correctness-preserving (cannot make correct code wrong), so this is safe. NOT YET VERIFIED on RDNA hardware — the only CDNA box available never executes the WMMA path. Validate with: ./test_sdpa_ref # on gfx1152/gfx1201 — L>=64 causal cases should pass
…it LDS)
The bwd kernel hardcoded BM=64 and the dispatch rejected D!=64, so head dims
above 64 (e.g. Qwen3.6 = 128) could not use the flash path. Forcing the kernel
at D=128/256 overflowed LDS (2*BM + BM*BN + 3*BM*D + 2*BN*D floats > 64KB) ->
garbage / NaN in test_sdpa_ref.
Make BM a template parameter and pick it per head dim so LDS fits a 64KB device:
D=64 -> BM=64 (60.5 KB)
D=128 -> BM=16 (41.1 KB) <- Qwen3.6
D=256 -> BM=8 (56.6 KB)
grid.y = ceil(qL/BM) and the dynamic-LDS byte count both derive from the chosen
BM, so shrinking BM only adds Q-tile blocks. The compute loops are D-agnostic
(for d < D), so this is a correctness-preserving retiling, not an algorithm
change. Support gate now allows D in {64,128,256}.
Pairs with the prior commit (missing dV/dS __syncthreads) which fixes the
wave32 race on the D=64 path.
Compile-verified for gfx942 (this f32acc kernel is plain-float, no WMMA
intrinsics). NOT yet numerically verified — validate on RDNA (gfx1152/gfx1201)
with ./test_sdpa_ref: all D in {64,128,256} x L in {8,64,256,1024} causal
should print OK.
…mmit) The previous commit changed sdpa_flash_bwd_smem(int) -> (int,int), but its callers in scaled_dot_product_attention.cpp (VJP::eval_gpu, VJP::use_fallback) still call the 1-arg form -> undefined reference at link. Keep the public signature int sdpa_flash_bwd_smem(int D). The BM-per-head-dim mapping is now a single constexpr helper sdpa_flash_bwd_bm(D) used by BOTH the size function and the kernel launch (BMv = sdpa_flash_bwd_bm(Dv)), so there is one source of truth and no drift. Compile-checked for gfx942.
Online softmax used if (tid < BLOCK_M) so wave0 lanes 16..31 read S rows written by wave1. rocWMMA store_matrix_sync to LDS is not made visible by plain __syncthreads(); that matched L<=16 OK and L>=32 diverge/NaN on gfx1152. Align softmax row ownership with the WMMA writer wave and flush with rocwmma::synchronize_workgroup after WMMA→LDS stores. Move wave_scale into dynamic smem and account for it in sdpa_flash_wmma_smem.
Device code used `#else WAVE=64` when __AMDGCN_WAVEFRONT_SIZE__ was unset.
On gfx115x that produced wave=tid/64 with a 128-thread RDNA launch, so only
rows 0..31 were written (L<=32 OK, L>=33 diverge/NaN). Use WARP_SIZE from
device/config.h (arch macros: RDNA→32, CDNA→64), matching host launch dims.
Verified on gfx1151: full D={64,128,256}×L={8..1024} causal OK (max|g-ref|<0.01).
Confirm exclusive 35B chat load passes; concurrent second process load SIGSEGV exit 139 after device bind (copy_contiguous/HIP path). Document ownership: robustness/OOM vs single-process correctness; no code fix PR yet.
CI uses pre-commit/action with --all-files. Format the tree so the docs investigation PR is not blocked by unrelated style drift on rocm-support sources. No intentional functional change.
docs(rocm): gfx1150 dual-process large-model load SIGSEGV (exclusive OK)
Did you ever get it working? |
|
@angeloskath @zcbenz How can I get this PR considered for Merge or what would be the easiest way for us to get it merged into mlx? I can provide AMD runners myself for consumer cards and hardware, I have plenty of equipment, RDNA3.5, 4 etc.... and a 10 GPU r9700 server I can offer for CI processes to do QA. I have been continually maintaining it and is implemented into an engine I am also maintaining, currently in alpha. I would really like all this hard work to be realized and pushed further. I believe if it got merged we could see a lot more contributions and help by the community. The AMD community is very strong and helpful. If it would help, I can squash the history grouped by person so all people who made changes and improvements keep their recognition while reducing the overall commit history. AMD also has released simulation software, where a CPU or another GPU can emulate another architecture of GPUs so we could have a full QA environment targeting all AMD GPUs so nothing is shipped without tests passing. |
Member
|
Currently we own all the code, i.e. we review and verify all code in this repo, so we can't just merge code that is maintained purely by community contributors, and we would require at least one full time team member to govern the new backend. Our team is hiring, but unfortunately we do not have a head count for maintaining a new backend, and we are already overloaded. I think the only reasonable way to move forward is to have AMD commit resources and developers to maintain the rocm backend in the long term, which is hard at the moment but should be realistic once we get ollama to use mlx-cuda. |
This was referenced Aug 10, 2026
21 tasks
This branch has not been deployed
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