fix(rocm): Gather reads narrow index dtypes as int64 and faults the GPU queue #1853

Part of #1801

Problem / Background

On ROCm, mlx::core::take (and every other Gather path) faults the GPU queue with HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION whenever the index array's dtype is narrower than 4 bytes. The faulting kernel names itself in the abort message as mlx::core::rocm::gather_general_kernel<float, long, 1>, that is, instantiated with IdxT = int64_t while the index buffer actually holds 1-byte or 2-byte elements, so each thread reads 8 bytes from a buffer sized for 1 or 2 and walks past the end. After the queue aborts the process does not die; it spins on the host, which is the symptom #1804 describes.

Current Behavior

Five dispatch sites in src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/indexing.hip choose the index type with a two-way branch of the form if (idx_dtype == int32 || idx_dtype == uint32) { ... int32_t ... } else { ... int64_t ... }:

Every integer dtype that is not int32 or uint32 therefore lands on the int64_t instantiation, including int8, uint8, int16 and uint16. MLX accepts any integer dtype for indices (dispatch_int_types in mlx/dtype_utils.h dispatches through MLX_INTERNAL_DTYPE_SWITCH_INTS, which covers all of them), so this is a fork-side defect, not an upstream one.

Measured boundary, on a Radeon 8060S (gfx1151), ROCm 10.0.0, at mlxcel f9ece5d9. Each row is take of a 16-entry float32 table with a 64-element index array, GPU against the CPU stream:

The index array's shape does not matter: [64], [8, 8], [1, 1, 64] and [1, 1, 1, 64] all fault with uint8 and all pass with int32.

Where mlxcel hits it: dequantize_v_turbo4 in src/lib/mlxcel-core/src/cache/turbo/quant.rs reaches ffi::take(&centroids_arr, &indices_u8, 0) on the 4-bit codebook (quant.rs:537, with the sibling call sites at quant.rs:618 and quant.rs:645), with uint8 indices. That is why every TurboQuant KV cache test faults once the Hadamard stub is gone (#1825): they previously aborted earlier, at Hadamard, and never reached this.

Minimal reproduction, as a standalone C++ program linked against the built MLX (no mlxcel involvement):

auto table = array(/* 16 floats */, {16}, float32);
auto idx = array(/* 64 uint8 values in [0, 16) */, {64}, uint8);
auto out = take(table, idx, 0, Device::gpu);
eval(out);  // HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION

Proposed Solution

The tradeoff is stated here rather than assumed, because measuring it is part of the work. The obvious change, giving each integer dtype its own instantiation, multiplies template instantiations by four at five sites, and gather_general_kernel is already templated on the value dtype (12 cases) and the index count (5 to 6 nidx branches per site), so the compile-time cost needs measuring before choosing it. The cheaper alternative is to normalize a narrow index array to int32 once on the stream before the launch, keeping the existing two instantiations, at the cost of one conversion pass over the index array. Either way the fix belongs in indexing.hip and must cover all five sites, not just the Gather fast-rows path, and it should be recorded in patches-rocm/LOCAL_FIXES.md.

Scope

In scope: all five dispatch sites in src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/indexing.hip (lines 694, 730, 896, 1039, 1186), a regression check covering the narrow index dtypes, and the patches-rocm/LOCAL_FIXES.md entry.

Out of scope: the Hadamard, FFT and SearchSorted implementations tracked in #1825, the error-surfacing work in #1804, and the fork-sync mechanics in #1813.

Implementation Notes

• Reuse: follow the existing macro structure in indexing.hip (ROWS_BY_T, DISPATCH_NIDX, DISPATCH_REDUCE, DISPATCH_IDX_TYPE) rather than introducing a parallel dispatch scheme, and record the fix in patches-rocm/LOCAL_FIXES.md in the same numbered form as the existing entries.
• Constraints: the fix lives in the ROCm patch tree only and must not change behavior on Metal or CUDA; if the normalize-to-int32 route is taken, the conversion must run on the same stream so it is ordered before the kernel launch.
• Edge cases: nidx == 0 (no index input, where idx_dtype defaults to int32), bool index arrays, multi-index gathers with mixed index dtypes across the index inputs, and index arrays large enough that int32 normalization would be lossy (an index count above INT32_MAX must still take the int64_t path).
• Error handling: an unsupported index dtype must throw a descriptive std::runtime_error naming the dtype, as the existing value-dtype default: arms do, rather than silently falling through to a mismatched instantiation.

Acceptance Criteria

• take with int8, uint8, int16 and uint16 indices matches the CPU stream on ROCm, at several index shapes (at minimum [64], [8, 8] and [1, 1, 1, 64]).
• The same holds for the scatter and gather-reduce paths that share the branch (Scatter, GatherAxis, ScatterAxis), not only Gather.
• The TurboQuant KV cache tests in mlxcel-core pass on ROCm, which is also an acceptance item of feat(rocm): implement FFT, Hadamard and SearchSorted on ROCm (audio, TurboQuant KV cache, stochastic speculative decoding abort today) #1825.
• A regression check exists so the branch cannot silently lose a dtype again.
• The change is recorded in patches-rocm/LOCAL_FIXES.md, and noted as a candidate to send back to the fork through chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813 since the defect is fork-side.
• The fix is integrated into the real ROCm code path that mlxcel-core calls, not left as a standalone helper.

Verification

cargo test -p mlxcel-core cache::turbo
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --all --check

A pass means the TurboQuant KV cache tests complete on the ROCm GPU stream with no HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION in the output and no host-side hang. Manual check on gfx1151: run the minimal reproduction above for each of the eight integer index dtypes and confirm all eight match the CPU stream.

Technical Considerations

The same class of defect was fixed once already in #1824, where Scatter passed an int32_t argument to an int64_t kernel parameter. The compile-time argument-size assertion added there (recorded as entry 14 in patches-rocm/LOCAL_FIXES.md) does not catch this one, because here the mismatch is between the kernel's template parameter and the dtype of the buffer it is pointed at, not between an argument and its parameter.