Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code
10. **mxfp4/mxfp8 in `gather_qmm`.** When a non-affine mode missed the affine-only fast paths in `GatherQMM::eval_gpu`, the generic dispatch instantiated `gather_qmv_kernel` with `AFFINE=true` and a T-typed scale pointer, so fp weights were dequantized with the affine formula (relative error 1.0 to 1e34 for mxfp8). Non-affine modes now launch `gather_qmv_kernel<T, uint8_t, BITS, 32, false>`; other fp group sizes throw. After the fix mxfp8 and mxfp4 `gather_qmm` match a dequantized f32 reference (relative error about 3e-3 for bf16) with sorted and unsorted indices. Applies to the fork; to be proposed there.
11. **`arg_reduce` launch geometry.** For short reduction axes `ArgReduce::eval_gpu` launched one 1024-thread block per output over a 1-D grid. AMD limits each grid dimension to 2^32 - 1 threads, so the launch failed with "invalid configuration argument" past about 4.2M outputs; mxfp4 `quantize` of a 4096x4096 matrix (an argmin over 16.7M elements) hit it. Axes of at most 64 elements now use a 64-thread block, and the grid spills into y. GPU `quantize` in mxfp4 is then bit-identical to the CPU at 4096x4096. Applies to the fork; to be proposed there.
12. **f16 in the fast `gather_qmm` path.** `GatherQMM::eval_gpu` sent only bf16 activations to `gather_qmv_warp_shared_kernel`; f16 fell through to the generic `gather_qmv_kernel`, which took about 205 ms per call on `gfx1151`. f16 checkpoints such as `mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit` decoded at about 0.1 tok/s, with 99.7% of GPU time in that kernel (rocprofv3). f16 now uses the same warp-shared kernel (the kernel was already templated on the element type), which takes 1.44 ms for a decode-shaped call at Mixtral's expert shape and matches a dequantized f32 reference. Applies to the fork; to be proposed there.
13. **f32 activations in the tiled qmv path.** `QuantizedMatmul::eval_gpu` entered the tiled qmv branch for any activation dtype, launched `qmv_tiled_kernel` only for bf16 and f16, and then returned, so f32 activations produced no kernel launch and the output was whatever the fresh allocation held (relative error from 1.0 up to 1e35 at every tested shape). The branch is now entered only for bf16/f16 and group sizes 32/64/128, so f32 takes the generic qmv dispatch, which matches a dequantized f32 reference (relative error about 1e-7). mlxcel-core's split/fused projection tests use f32 inputs and exposed it. Applies to the fork; to be proposed there.

## Dropped from the vendored set

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3835,9 +3835,15 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
// Restrict to 4/8-bit; 2/5/6-bit fall through to the warp-shared QMV kernel.
bool tiled_bits_supported = (bits_ == 4 || bits_ == 8);
// tile_n can be 4 on low-CU gfx1152 — still a valid tiled path.
// mlxcelverse: the launches below cover only bf16/f16 and group sizes
// 32/64/128, but the branch returns unconditionally, so any other input
// (f32 activations in particular) left the output unwritten and returned
// whatever the allocation held. Only enter when a kernel will launch.
if (use_tiled && tiled_bits_supported && use_fast_qmv &&
!can_use_batched_qmv && tile_n >= 4 &&
mode_ == QuantizationMode::Affine) {
mode_ == QuantizationMode::Affine &&
(x.dtype() == bfloat16 || x.dtype() == float16) &&
(group_size_ == 32 || group_size_ == 64 || group_size_ == 128)) {
{
dim3 tiled_block(host_warp, tile_n);
const int n_tiles = (N + tile_n - 1) / tile_n;
Expand Down