Repository navigation
perf(rocm): route f16 gather_qmm to the warp-shared kernel - #1822
Merged
Merged
Conversation
GatherQMM on ROCm sent only bf16 activations to gather_qmv_warp_shared_kernel. f16 checkpoints fell through to the generic gather_qmv_kernel, which takes about 205 ms per call on gfx1151: mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit (an f16 checkpoint) decoded at about 0.1 tok/s, and rocprofv3 attributed 99.7% of GPU time to that one kernel. Qwen3-30B-A3B was unaffected because it is bf16. The warp-shared kernel is already templated on the element type, so f16 now takes it too. A decode-shaped call at Mixtral's expert shape (E8, N14336, K4096, 2 rows) drops from about 205 ms to 1.44 ms and still matches a dequantized f32 reference (relative error about 4e-4). The fix is recorded as item 12 in patches-rocm/LOCAL_FIXES.md. The sorted large-row path used by prefill stays slow for both dtypes because the expert-batched kernel is opt-in after its bf16 correctness bug; that is left to #1814. Refs #1814
inureyes
added
status:review
This was referenced Sep 11, 2026
inureyes
added a commit
that referenced
this pull request
Sep 12, 2026
The experimental ROCm backend had no evidence that it decodes the same text as Metal. This adds that evidence, the gate that produced it, and a build fix the gate exposed. Twelve model x width pairs (qwen3-0.6b, llama-3.1-8b-instruct, qwen3-30b-a3b, mixtral-8x7b-instruct at widths 1, 8 behind 512 tokens of context, and 256) traced on a Radeon 8060S (gfx1151) at bec6474 against Metal traces from the same commit, corpus, arguments and checkpoint revisions. The two backends disagree on the top-1 token at 114 of 5120 positions, and the largest reference top-two gap at which they ever disagree is 1.125 logits: wherever the reference preferred its own choice by more than that, ROCm made the same choice. 100 of the 114 disagreements take the reference's rank-2 token. A same-backend control (swapping the fused MoE kernel for gather_qmm within Metal) moves only one position, at a gap of 0.062. docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md has the matrix, a threshold sweep at 0.5 / 1.0 / 2.0 so the --decided line is inspectable rather than asserted, the rank distribution, and why byte-identity is the right criterion where it holds but not across backends. The Mixtral ROCm traces come from bec6474 plus only the f16 gather_qmm dispatch merged as #1822, which lives entirely in src/lib/mlx-cpp/patches-rocm and touches no Rust or bridge code; without it that checkpoint spends about 205 ms per MoE gather_qmm call. Every other ROCm trace and all of the Metal side are plain bec6474. make verify-test-rocm and make test-fast-rocm mirror the CUDA targets. That gate, not the traces, found the defects fixed in #1823 and #1824 plus the one fixed here: mlxcel-surgery and mlxcel-xla test binaries never received the ROCm rpath, because a dependency's cargo:rustc-link-arg does not reach the crate that links it, so they built and then failed to start with "libamdhip64.so.7: cannot open shared object file". The emission now lives in build_support/rocm_rpath.rs and all three crates call it. The gate goes from 12 failing targets to 5, all of them #1825 and #1803. compare_logit_traces.py now prints the largest gap at a disagreement, so the threshold-free bound is reproducible from any trace pair, and it reports inconclusive rather than a pass when the reference has no decided position at all. Refs #1809, #1802
1 of 5 tasks
This was referenced Sep 29, 2026
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.
On ROCm,
GatherQMMsent only bf16 activations togather_qmv_warp_shared_kernel. f16 checkpoints fell through to the genericgather_qmv_kernel, which takes about 205 ms per call ongfx1151, somlx-community/Mixtral-8x7B-Instruct-v0.1-4bit(an f16 checkpoint) decoded at about 0.1 tok/s. rocprofv3 put 99.7% of GPU time in that one kernel. bf16 models such as Qwen3-30B-A3B were not affected.The warp-shared kernel is already templated on the element type, so f16 now takes it too. The change is one dispatch block in the ROCm overlay (
src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip), recorded as item 12 inLOCAL_FIXES.md. Metal and CUDA builds never copy this directory.Measured on
gfx1151:gather_qmmat Mixtral's expert shape (E8, N14336, K4096, 2 rows, 4-bit g64): about 205 ms to 1.44 ms, matching a dequantized f32 reference (relative error about 4e-4); bf16 is unchanged.mlxcel generateon Mixtral-8x7B-Instruct-4bit withMLXCEL_FUSED_MOE=0: 7.45 tok/s for 60 tokens, with coherent output, against about 0.1 tok/s before.Not addressed: the sorted large-row path that prefill uses stays slow for both dtypes (hundreds of ms per call at this shape) because the expert-batched kernel is opt-in after its bf16 correctness bug. That belongs to #1814, which this PR does not close.
Refs #1814