Conversation
load_vector and load_vector_safe add the inputs of each sub-byte chunk in the input type before widening the sum. With bfloat16 or float16 inputs that intermediate rounds, so the bias contribution bias * sum(x) is off even when the exact sum is representable in float. Widen each input before adding it, for every sub-byte width (2, 3, 4, 5 and 6 bits); the 8-bit branch already accumulates one value at a time in float. quantized_nax.h carries unused copies of both functions and gets the same change so the two headers stay in sync. The test builds inputs whose running sum is not representable in the half precision types and checks the exact result for several output sizes, input sizes, group sizes, bit widths and dtypes.
When the input size is aligned for qmv_fast but the output size is not a multiple of 8, qmv falls back to the generic kernel. For outputs below 8 rows its branch bounds-checks every row inside the reduction loop, which makes each call several times slower on the GPU than the fast kernel. For outputs of 8 rows or more that are not a multiple of 8 the generic kernel takes its other branch, which moves the last tile back over rows it has already computed and so needs no per-row check either. That branch is still the slower one: it reads a narrower block per thread than qmv_fast for every width but 2 bits, and it keeps the guarded tail that an aligned input does not need. Give qmv_fast_impl a partial_rows flag: rows of the last SIMD group that fall past the output reuse the weights of the last valid output row and are not stored, so the reduction loop stays free of per-row checks. qmv uses the new affine_qmv_fast_rows kernel for aligned affine inputs whose output size is not a multiple of 8; other routes are unchanged.
wyanzhao
marked this pull request as ready for review
September 15, 2026 08:33
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
2 participants
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.
Fixes #4498.
Affine QMV currently falls back to the generic kernel when the output size is not divisible by 8, even when the input meets the
qmv_fastalignment requirement. This is particularly expensive for narrow, dependent projections such as the N=1 and N=4 calls in autoregressive decode.This PR adds
affine_qmv_fast_rows. In the final SIMD group, inactive lanes reuse the last valid weight row and skip their stores, so the reduction loop can stay the same asqmv_fast. The new path is used whenqmvis selected for an aligned affine call without a global scale. Gathered QMV and the floating-point quantization modes keep their existing kernels. The extra instantiations add about 1 MB tomlx.metallib.While testing the new path, I found that the affine bias term accumulated
sum(x)in the input type before widening it. That loses small contributions for float16 and bfloat16 inputs. The first commit widens each value before accumulation for 2- through 6-bit quantization; the 8-bit path already did this.Performance
Measured on an Apple M5 Max with bfloat16 inputs. Baseline uses the generic kernel and Candidate uses
affine_qmv_fast_rows; both include the precision fix. The ratio is the geometric mean of paired Baseline/Candidate time ratios from A/A-calibrated runs, so values above 1 favor Candidate.For the model-shaped benchmark, I used layers 4–7, one complete four-layer cycle with three GDN layers followed by one QSA layer. The benchmark configuration cannot run the full model, so I did not measure whole-model generation.
Correctness
test_qmv_affine_bias_sum_precisionfails on main for the affected float16 and bfloat16 cells below 8 bits and passes on this branch.test_qmv_fast_rowscovers output sizes 1, 3, 5, 12 and 17 across 2- through 8-bit affine quantization, including batched weights. Its M=3 cases useqmv_wideon M3 and newer GPUs, so I could not exercise theiraffine_qmv_fast_rowsroute on this machine.For the layer benchmark, Baseline and Candidate produced bitwise-identical outputs over eight decode steps at both context lengths. Every changed call was also within half a bfloat16 step of a float64 reference.
Validation and quick reproduction
Both new tests also pass with
MLX_METAL_JIT=ON. The fulltest_quantizedrun has the same 1273 failing subtests on main and this branch; 1188 are intest_qmv_wide.The reproducer in #4498 shows the original output-size cliff. To reproduce the N=1 row above, set
IN_DIM = 2560,GROUP_SIZE = 64,BITS = 8, and includeout_dim = 1, then run it once against main and once against this branch. It prints differential median time rather than the paired ratio in the table, and results depend on the machine and build.The layer benchmark uses
Qwen4ExpDecoderLayerlayers 4–7 from mlx-vlm1ecf1ecdd28af102eded679be0daa5c76ab2a068and tensors frommlx-community/Qwen3.8-Flash-Next-4bitrevision07b5dc6c54600a359b87f1e53e7adf6351c72a2c. It restores a cache prefetched with 128 or 4096 synthetic tokens before each sample and times 64 decode steps throughQwen4ExpBatchInvariantForward._layer.