Repository navigation
feat(rocm): implement Hadamard and SearchSorted, fix narrow index dtypes - #1856
Merged
Merged
Conversation
The ROCm backend left `SearchSorted` a `NO_GPU` stub, and the throw crosses the cxx bridge into a `noexcept` extern, so every hit ended in `std::terminate`. 13 tests in `speculative::distribution_tests` and `speculative::stochastic_accept` aborted the mlxcel-core test binary because of it. mlxcel has no `searchsorted` call site of its own; TurboQuant's nearest-centroid lookup is a host-side binary search. It is reached from MLX internals on the sampling path: `residual_resample` calls `fused_sample`, which falls back to `random::categorical` on a backend with no Gumbel-max kernel, and `categorical_impl` takes its inverse-CDF branch for 1-D logits, which is mlxcel's single-row sampling case. Landing the ROCm Gumbel-max kernel under #1803 would remove the dependency from the fused path, though the graph fallback keeps it. The kernel is ported from `mlx/backend/cuda/sort.cu` into `mlx/backend/rocm/sort.hip` and reuses this backend's `LessThan`, so NaN ordering matches the CPU without a second comparison. It differs from the CUDA original in being grid-stride: this backend caps a launch grid, so one index per thread would silently leave the tail of a large output unwritten, which is the same class of defect as #1823. Verified against the CPU stream on both the left and right sides for duplicate runs, strictly increasing sequences, out-of-range queries, trailing NaN, a single element, and 17M queries, which is past the 65535 x 256 grid extent and so exercises the grid-stride path. The 29 previously aborting tests now pass: `speculative::stochastic_accept` 18 of 18, `speculative::distribution_tests` 16 of 16. Refs #1825
The ROCm port of the Hadamard transform in #1825 needs a line to pass, and the two things worth checking are not the same: whether the transform is scaled the way the TurboQuant KV cache assumes, and how much a half-precision round trip drifts. The probe reports both per head dim and dtype. The norm ratio catches a scale mistake on its own, which the round trip cannot: an unscaled Hadamard is not an involution either, but it reads as sqrt(N) immediately. The round-trip error gives the band a port has to stay inside. Measured on an M1 Ultra at mlxcel `1f107710` with MLX pin `81ba1c6a`: the norm ratio is 1.000000 on every row in both dtypes, confirming the default 1/sqrt(N); fp16 round-trip relative RMS runs 3.8e-4 to 5.0e-4 with max absolute 1.953e-3, one row at 3.906e-3; f32 relative RMS is about 1e-7 with max absolute 3.0e-7 to 7.2e-7. The header records those numbers and the pass bands derived from them. Refs #1825, #1809
`Hadamard` was a `NO_GPU` stub on ROCm, and the throw crosses the cxx bridge into a `noexcept` extern, so it aborted the process. That took out every TurboQuant KV cache mode and `tests/wht_op.rs`.
`mlx/backend/rocm/hadamard.hip` adds a single-stage radix-2 butterfly in shared memory, covering `m == 1` with a power-of-two last axis up to 8192. That is the whole of what 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, so the sizes that arrive are head dims, chiefly 64, 128 and 256. Larger `n` and the `m` in {12, 20, 28} families throw a message naming the limit instead of returning a wrong answer; covering them needs the radix-m codelet the CUDA path generates through NVRTC.
The scale is the primitive's `scale_`, which MLX defaults to `1/sqrt(n)`. That default is load-bearing: it makes the transform an involution, and mlxcel's TurboQuant rotation relies on `wht(wht(x))` returning `x`. A kernel with a hardcoded scale of 1 or `1/n` would pass a round-trip test that only checks shape or finiteness while silently corrupting KV values, so the forward norm ratio is measured rather than assumed. Accumulation is in float whatever the element type.
Like the CUDA original the grid is clamped at 65535 blocks and the kernel is grid-stride; this backend's `get_launch_args` clamps without that pairing, which is why it is not used here.
Verified with `examples/wht_numeric_probe` (cherry-picked from the Metal side so both backends run the same code) against the M1 Ultra reference at head_dim 64, 128 and 256: norm ratio 1.000000 within 1e-5, which is what separates the orthonormal scale from an unscaled transform's sqrt(N); fp16 relative RMS about 2.0e-4 against Metal's 3.8e-4 to 5.0e-4; f32 relative RMS about 1e-7, matching Metal. `tests/wht_op.rs` passes 5 of 5 on gfx1151.
The TurboQuant cache tests still do not pass. With the stub gone they reach a separate pre-existing defect: `Gather` on ROCm reads an index array narrower than 4 bytes as `int64_t`, which faults. That is filed separately.
Refs #1825
Five dispatch sites in `mlx/backend/rocm/indexing.hip` chose 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 that reads eight bytes per element out of a one or two byte buffer. On ROCm that is an `HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION`: the GPU queue aborts, and the process then spins on the host rather than failing, which is the symptom #1804 describes. MLX accepts any integer dtype for indices. mlxcel reaches it with uint8, because the TurboQuant 4-bit codebook lookup is a `take` with uint8 indices. That is why every TurboQuant KV cache test faulted as soon as the `Hadamard` stub was gone: they used to abort earlier and never got here. Narrow index arrays are now widened to int32 once, before the dispatch, in `Gather`, `Scatter`, `GatherAxis` and `ScatterAxis`. The obvious alternative, giving each of the eight integer dtypes its own instantiation, was implemented and measured first. It is also correct, and it is expensive: the index type multiplies an existing product of twelve value dtypes and five to six index-count branches, and compiling this translation unit went from 26 s to 348 s. Widening keeps the two instantiations and costs one copy pass over the index array, on the paths that are broken today and nowhere else. Verified against the CPU stream with a standalone program linked against the built MLX, for int8, uint8, int16, uint16, int32, uint32, int64 and uint64 at index shapes `[64]`, `[8, 8]`, `[1, 1, 64]` and `[1, 1, 1, 64]`. On gfx1151 the previously faulting test groups now pass: `cache::turbo` 164 of 164, `cache::sparse_v` 19 of 19, `cache::detach` 51 of 51. The defect is fork-side, not upstream, so it is a candidate to send back through #1813. The compile-time argument-size assertion added in #1824 does not catch this class: there the mismatch was between an argument and its kernel parameter, here it is between a kernel's template parameter and the dtype of the buffer it is pointed at. Closes #1853
Two lines in the cherry-picked probe were over the width rustfmt enforces, so `cargo fmt --all -- --check` failed on this branch. No behavior change. Refs #1825
inureyes
added
type:enhancement
The file comment said the 1/sqrt(n) default makes the transform an involution, which is true and reads as if the round trip were a property of the transform. It is a property of that one scale: H * H = n * I, so applying scale s twice gives s^2 * n * I. A caller passing its own scale, which the quantization work in #1807 and #1808 may want in order to fold the rotation into a block scale, needs the inverse scale 1/(s * n) and the forward norm ratio s * sqrt(n) instead of a round trip. Comment only. Refs #1825
5 tasks
This was referenced Sep 29, 2026
Closed
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.
Two of the three missing ROCm primitives from #1825, plus the defect the second one exposed.
FFTis 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.SearchSortedis ported frommlx/backend/cuda/sort.cuintosort.hip, reusing this backend'sLessThanso NaN ordering matches the CPU. mlxcel has no call site of its own: it is reached fromresidual_resample, which callsfused_sample, which falls back torandom::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.Hadamardis a shared-memory radix-2 butterfly in a newhadamard.hip, coveringm == 1with 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. Largernand themin {12, 20, 28} families throw a message naming the limit instead of returning a wrong answer. The scale is the primitive'sscale_, which MLX defaults to1/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
Hadamardstub gone, the TurboQuant tests reached a separate pre-existing defect, filed as #1853 and fixed here. Five dispatch sites inindexing.hippicked the kernel's index type withint32_tfor int32 and uint32 andint64_tfor 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: anHSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATIONthat aborts the GPU queue and leaves the process spinning. mlxcel hits it through the TurboQuant 4-bit codebook lookup, atakewith uint8 indices.Both fixes for that were implemented and measured before choosing. Instantiating every integer dtype is correct and takes
indexing.hipfrom 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) atf9ece5d9:examples/wht_numeric_probe, head_dim 64 / 128 / 256tests/wht_op.rsspeculative::stochastic_accept,speculative::distribution_testscache::turbo,cache::sparse_v,cache::detachtakeagainst the CPU stream, 8 index dtypes x 4 shapessearchsortedagainst the CPU stream, both sidesmake verify-test-rocmgoes from 5 failing targets to 4, and the whole remainder belongs to other issues: themlxcellib is 8200 passed with 1 failure and 3 aborts (FFT, and one[cuda_kernel]path from #1803), andmlxcel-coreis 1701 passed with 2 failures from #1805 and 36 aborts that are all #1803 apart from two nvfp4 group-16 cases from #1806. NoHadamardorSearchSortedabort is left.examples/wht_numeric_probe.rsis 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 inpatches-rocm/.Closes #1853. Refs #1825.