Refresh MLX fork onto current upstream cutoff - #3
Merged
Conversation
…mory (ml-explore#3872) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
Co-authored-by: Anastasiia Filippova <a_filippova@apple.com>
Co-authored-by: Anastasiia Filippova <a_filippova@apple.com>
…-explore#3914) Co-authored-by: Anastasiia Filippova <a_filippova@apple.com>
Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
…ore#4222) Co-authored-by: codeAnqiang-ma <273298913+codeAnqiang-ma@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com>
…-explore#4125) Co-authored-by: Cheng <git@zcbenz.com>
affine qmv_fast set packs_per_thread=2 for all bits except 2-bit, so 1-bit got values_per_thread=64 (x_thread[64], ~256B/thread) -> low occupancy -> 1-bit decode saturates only ~75% of M5 DRAM BW vs ~90/96% for 2/4-bit. Use 1 pack/thread for bits<=2 (values_per_thread=32), matching 2-bit's register footprint. Measured on M5 Pro (distinct-weight DRAM-bound, 2-bit as drift control): 1-bit 24.0 -> 21.9 us/matvec (~9%, 75->82% BW), 2-bit control 33.2->32.9 (0.8% drift); correct, rel_err 2.7e-4. Also makes scale_step_per_thread (=group_size/values_per_thread) well-defined for group_size=32 at 1-bit.
The qmv_wide kernel dequantizes a full group of weights into registers and reuses them across the M input rows. That amortization only pays off for 2-bit once three or more rows share a group; at M=2 it breaks even, and for 1-bit the weight traffic is small enough that the per-row dequant dominates and the specialized qmv is faster. Gate the dispatch accordingly: 1-bit and 2-bit M<2 keep the specialized qmv; 2-bit routes to qmv_wide only at M>=3 on gen-15+; fp modes are unchanged. Add 1-bit affine coverage (full sweep + tiny shapes) to test_qmv_wide. Measured on affine 2-bit matvecs across a range of projection shapes: qmv_wide wins at widths >=3 (more so as N grows), with no regression at width <=2 or for 1-bit.
Adapt the bit-exact NVFP4 loader rewrite from mlx.fast submission 4228f9e1-18a5-4a1a-8047-96dc4e9289ca (validated commit b8fe3af). Fold fp4's exact 2^14 renormalization into the e4m3 group scale and decode eight packed nibbles per uint32 in the regular and NAX loaders. Other quantization modes keep the scalar path.\n\nThe source submission exhaustively checked 404,226,048 staged values with zero bit mismatches. This adaptation adds a local 524,288-value BF16 equivalence gate in the parent mlx-swift repository.\n\nCo-authored-by: anupsv <6407789+anupsv@users.noreply.github.com>
This was referenced Aug 14, 2026
a-github-name
marked this pull request as ready for review
August 14, 2026 13:51
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
20 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.
Summary
bd5c3a2b170bb95340482e35b2a49fb08aea4de3Final revisions
b57bd7640f3f7c743b76a58478faaf1e8ee084f2bd5c3a2b170bb95340482e35b2a49fb08aea4de3Validation
uvx pre-commit run --all-filesb57bd7640f3f7c743b76a58478faaf1e8ee084f2Linux/CUDA proof
The exact core and Swift revisions were built and executed on Ubuntu 24.04 arm64 with an NVIDIA GB10 (compute capability 12.1), driver 580.159.03, CUDA 13.3, Swift 6.0.3, cuDNN, NCCL, and OpenBLAS.
scripts/prepare-linux-native.shbuilt all 246 MLX/Swift bridge targets, and MLX Swiftexample1 --device gpuran successfully onDevice(gpu, 0). The downstream full release CUDA package linked and launched with toolkit environment variables unset, resolved its CUDA/cuDNN/NCCL/NVRTC/cuFFT libraries, and selected native MLX GPU. The final mere.run Linux CPU, macOS Swift, app-bundle, docs, and security checks are green on the exact immutable pins.Downstream performance A/B
The dependent mere.run refresh was release-benchmarked against the previous production pin on an M4 Max. Laguna XS 2.1 remained flat at
107.76 -> 107.75 tok/sdecode with equal0.750 smedian prefill and byte-identical 128-token output. MiniMax-H3 remained flat at17.45 -> 17.51 smedian wall time, with its steady 50-block transformer evaluation0.828 -> 0.827 sand unchanged Metal memory. Full receipts are recorded on sawfwair/mere-run#297.