Repository navigation
fix(rocm): stop the tiled qmv path from skipping f32 activations - #1823
Merged
Merged
Conversation
inureyes
added
status:review
On ROCm, QuantizedMatmul entered the tiled qmv branch for any activation dtype, launched qmv_tiled_kernel only for bf16 and f16, and then returned. With f32 activations no kernel ran, so quantized_matmul returned the contents of a fresh allocation: relative error from 1.0 up to 1e35 against a dequantized reference at every shape tested, from 16x64 to 4096x4096. bf16 and f16 were correct, which is why the model-level checks did not see it. The mlxcel-core split/fused projection tests (layers::tests::*) use f32 inputs and failed on the ROCm test gate. The branch is now entered only for the dtypes and group sizes it launches (bf16/f16, group 32/64/128), so f32 falls through to the generic qmv dispatch. Recorded as item 13 in patches-rocm/LOCAL_FIXES.md. Validated on gfx1151: f32 quantized_matmul matches a dequantized f32 reference (about 1e-7) for N x K from 16x64 to 4096x4096 and M = 1, 8; a sweep over f32/f16/bf16, K 64-256, N 4-24 and 4/8 bits has no failure; affine bf16/f16 checks are unchanged. Refs #1809
inureyes
force-pushed
the
fix/issue-1809-rocm-f32-tiled-qmv
branch
from
September 12, 2026 00:34
12ff5b0 to
9842d4a
Compare
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
This was referenced Sep 12, 2026
inureyes
added a commit
that referenced
this pull request
Sep 12, 2026
Two of the three missing ROCm primitives from #1825, plus the defect the second one exposed. `FFT` is not here; it needs hipFFT linked and a plan cache, and its blast radius turned out to be two call sites rather than the whole audio stack. `SearchSorted` is ported from `mlx/backend/cuda/sort.cu` into `sort.hip`, reusing this backend's `LessThan` so NaN ordering matches the CPU. mlxcel has no call site of its own: it is reached from `residual_resample`, which calls `fused_sample`, which falls back to `random::categorical`, whose inverse-CDF branch for 1-D logits is mlxcel's single-row sampling case. Unlike the CUDA original the kernel is grid-stride, because this backend caps a launch grid and one index per thread would silently leave the tail of a large output unwritten, the same class of defect as #1823. `Hadamard` is a shared-memory radix-2 butterfly in a new `hadamard.hip`, covering `m == 1` with a power-of-two last axis up to 8192. That is all mlxcel reaches: `wht()` asserts a power-of-two last axis and the KV cache advisor demotes Turbo modes for models whose head_dim is not one. Larger `n` and the `m` in {12, 20, 28} families throw a message naming the limit instead of returning a wrong answer. The scale is the primitive's `scale_`, which MLX defaults to `1/sqrt(n)`; that default is load-bearing, because it makes the transform the involution TurboQuant's rotation relies on, so the forward norm ratio is measured rather than assumed. With the `Hadamard` stub gone, the TurboQuant tests reached a separate pre-existing defect, filed as #1853 and fixed here. Five dispatch sites in `indexing.hip` picked the kernel's index type with `int32_t` for int32 and uint32 and `int64_t` for everything else, so int8, uint8, int16, uint16 and bool index arrays reached a kernel reading eight bytes per element from a one or two byte buffer: an `HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION` that aborts the GPU queue and leaves the process spinning. mlxcel hits it through the TurboQuant 4-bit codebook lookup, a `take` with uint8 indices. Both fixes for that were implemented and measured before choosing. Instantiating every integer dtype is correct and takes `indexing.hip` from 26 s to 348 s to compile, because the index type multiplies twelve value dtypes and five to six index-count branches; upstream does instantiate per dtype in both Metal and CUDA but compiles those kernels at runtime, so the combinations are free there and not here. Widening a narrow index array to int32 once keeps the two instantiations and costs one copy pass, only on paths that are broken today, so that is what landed. Numbers, on a Radeon 8060S (`gfx1151`) at `f9ece5d9`: | Check | Result | |---|---| | `examples/wht_numeric_probe`, head_dim 64 / 128 / 256 | norm ratio 1.000000 within 1e-5; fp16 relative RMS about 2.0e-4 against the M1 Ultra reference's 3.8e-4 to 5.0e-4; f32 about 1e-7, matching | | `tests/wht_op.rs` | 5 of 5 | | `speculative::stochastic_accept`, `speculative::distribution_tests` | 18 of 18 and 16 of 16, previously 29 aborts | | `cache::turbo`, `cache::sparse_v`, `cache::detach` | 164, 19 and 51, all previously aborting | | `take` against the CPU stream, 8 index dtypes x 4 shapes | all match | | `searchsorted` against the CPU stream, both sides | matches on duplicates, strict increase, out-of-range, NaN, single element, and 17M queries past the grid extent | `make verify-test-rocm` goes from 5 failing targets to 4, and the whole remainder belongs to other issues: the `mlxcel` lib is 8200 passed with 1 failure and 3 aborts (`FFT`, and one `[cuda_kernel]` path from #1803), and `mlxcel-core` is 1701 passed with 2 failures from #1805 and 36 aborts that are all #1803 apart from two nvfp4 group-16 cases from #1806. No `Hadamard` or `SearchSorted` abort is left. `examples/wht_numeric_probe.rs` is cherry-picked from the Metal side so both backends measure with the same code rather than trading numbers; its hardcoded host name is removed, since it mislabels every run elsewhere. Validated: `cargo fmt --all -- --check`, `cargo clippy --workspace --all-targets --features rocm -- -D warnings`, `scripts/ci/check_cross_repo_refs.py`, `scripts/insert_apache_header.py --check`. Metal and CUDA builds copy no file in `patches-rocm/`. Closes #1853. Refs #1825.
5 tasks
This was referenced Sep 29, 2026
fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback
#1808
Closed
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,
QuantizedMatmulentered the tiled qmv branch for any activation dtype, launchedqmv_tiled_kernelonly for bf16 and f16, and then returned. With f32 activations no kernel ran, soquantized_matmulreturned the contents of a fresh allocation: relative error from 1.0 up to 1e35 against a dequantized reference at every shape tested (16x64 to 4096x4096). bf16 and f16 were correct, which is why the model-level checks did not see it. The ROCm test gate for #1809 did: the mlxcel-core split/fused projection tests use f32 inputs.The branch is now entered only for the dtypes and group sizes it actually launches, so f32 takes the generic qmv dispatch. The change is one condition in the ROCm overlay (
src/lib/mlx-cpp/patches-rocm/.../quantized/qmm.hip), recorded as item 13 inLOCAL_FIXES.md; Metal and CUDA builds never copy this directory.Validated on
gfx1151:quantized_matmulmatches a dequantized f32 reference (about 1e-7) for N x K from 16x64 to 4096x4096 and M = 1, 8, and a sweep over f32/f16/bf16, K 64-256, N 4-24, 4/8 bits has no failure; bf16/f16 affine checks are unchanged.layers::tests::{mixed_bit_width_qkv_forward_matches_independent_projections, split_and_fused_projections_agree_on_uniform_weights, split_project_concat_reproduces_the_fused_activation}andffi_tests::compiled_qgelu_mlp_forward_*.This touches the same file as #1822 in a different hunk;
LOCAL_FIXES.mdnumbers this fix 13 so the two merge in either order.Refs #1809