Skip to content

chore(deps): bump the MLX pin and declare a base commit for each overlay #1932

Description

@inureyes

Problem / Background

Upstream ml-explore/mlx main is ahead of our pin 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a, and the range includes a quantized-matmul correctness fix on a file we overlay. Moving the pin is a one-line edit, but it is only safe if we know which upstream commit each overlay was copied from: overlays are applied whole-file via configure_file(... COPYONLY), so any upstream change to an overlaid file is silently reverted the moment the pin moves. Today only 6 of 30 overlays record a usable base commit, so the audit is a manual read of the upstream diff every time. #1042 did that audit by hand for 21 overlays; #824 / #829 shipped a pre-fix rms_norm.cu overlay that reverted an upstream CUDA fix and needed #830 / #831 to undo.

The bump and the base declarations belong in one PR: the declared bases are what make this bump checkable, both live in src/lib/mlx-cpp/patches*, and one verification pass covers both.

Current State (measured 2026-09-19)

  • Pin 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a, written only at src/lib/mlx-cpp/CMakeLists.txt:171 (the GIT_TAG line inside FetchContent_Declare(mlx ...)), single source of truth since chore(build): make the MLX pin a single source of truth so a partial bump cannot silently build the old MLX #1047. scripts/ci/mlx_pinned_commit.sh prints it; src/lib/mlxcel-core/build_support/mlx_pin.rs parses it for build.rs.
  • Upstream main was 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 (upstream PR #4507) when measured, 14 commits ahead of the pin. Re-measure at implementation time; the target is upstream main at that moment.
  • 30 overlay files: 27 under src/lib/mlx-cpp/patches/, 3 under src/lib/mlx-cpp/patches-cuda/. Copied by configure_file(... COPYONLY) at src/lib/mlx-cpp/CMakeLists.txt:45 (Metal list), :61 (CUDA glob) and :78 (patches-cuda glob, CUDA builds only).
  • 6 declare a base commit: patches/mlx/backend/cuda/cuda_utils.h (81ba1c6a), patches/mlx/backend/cuda/device/qmm_sm80.cuh (b7c3dd6d), patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu (b7c3dd6d), patches/mlx/backend/cuda/binary/binary.cuh (e9463bb), patches-cuda/fast.cpp (81ba1c6a), patches-cuda/ops.cpp (81ba1c6a). The drift check below returns 0 for all six against the current pin.
  • The other 24 have no usable base: 6 name a release tag rather than a commit (v0.31.1, in the five patches/mlx/backend/cuda/reduce/* files and patches-cuda/dtype.cpp), and 18 declare nothing. Two of those 18, patches/mlx/backend/cuda/quantized/quantized.cpp and patches/mlx/backend/metal/quantized.cpp, mention "Synced to upstream 81ba1c6a" in prose, which no check can read.

The four overlays that conflict with this bump

Between 81ba1c6a and 59d600b5e, upstream changed four files we overlay, all of them mlxcel's own custom quantized kernels, so each is a rebase and not a re-copy:

  • patches/mlx/backend/cuda/device/qmm_naive.cuh
  • patches/mlx/backend/cuda/quantized/qmm/qmv.cu
  • patches/mlx/backend/cuda/quantized/quantized.cpp
  • patches/mlx/backend/metal/quantized.cpp

The upstream commits in range for those files are dfe17bafb (upstream #3912, "Fix fp quantized matmul corruption when the quantized dim is not a multiple of 32"), 3eb822cac (upstream #4483, "Load global scales in qmm_t kernels"), d9add9d11 (upstream #4481) and 59d600b5e (upstream #4507, "[CUDA] Add global scale support to gather_qmm"). Re-derive the per-file mapping with git log --oneline 81ba1c6a..<new pin> -- <path>. dfe17bafb is the one to watch: it is a correctness fix, and a naive bump reverts it.

Proposed Solution

Four phases in one PR, in this order, because Phase 1 is what tells Phase 2 which files conflict and Phase 3 is only reviewable once 1 and 2 are on the same branch.

Phase 1, declare a base for all 30 overlays, against the current pin, before touching it. Follow the header convention that already exists in patches/mlx/backend/cuda/cuda_utils.h and patches/mlx/backend/cuda/device/qmm_sm80.cuh: after the copyright line, a comment block saying what the patch does, then exactly one line // Modified from upstream MLX <sha> <upstream/relative/path>, then the change in one or two lines. The sha may be abbreviated as long as git rev-parse resolves it in an MLX clone. Record the base per file, not in a shared UPSTREAM file: patches-rocm/UPSTREAM works because that whole overlay was retargeted at one commit, whereas these 30 were copied at different times (bases seen today span v0.31.1, e9463bb, b7c3dd6d, 81ba1c6a). Determine each real copy point instead of stamping the pin everywhere: diff the overlay against git show <candidate>:<path> and pick the newest commit where only mlxcel's intended delta remains, walking back through git log --oneline -- <path> until it does. Resolve the v0.31.1 bases with git rev-list -n 1 v0.31.1. If a file turns out to lag upstream (a pure revert), say so in the header and either refresh it or file a follow-up.

Phase 2, rebase the four conflicting overlays. Mirror src/lib/mlx-cpp/patches-rocm/README.md: git merge-file <overlay> <file at old base> <file at new pin>, then compare the result line by line against the new upstream file. Every mlxcel delta must survive; each is documented in the overlay's own header, and metal/quantized.cpp's is also in the CMakeLists comment at src/lib/mlx-cpp/CMakeLists.txt:31-35 (the MLXCEL_QMV_WIDE off-switch on use_qmv_wide, three hunks, #1186 / #1187). The others are quantized/quantized.cpp (input contiguity, 65535 grid chunking #648, sorted M==1 GatherQMM grouped-GEMM route #629, MLXCEL_TRACE_ARCH one-shot trace #1537) and qmv.cu (multirow qmv #725, row-window cap #906, pre-Ampere accumulator #1539). Update each rebased file's base line to the new pin.

Phase 3, bump the pin. Edit the single GIT_TAG line, keeping the shape both parsers require (documented in the comment block directly above it): full 40-character lowercase hex, first token on its line, inside that declaration.

Phase 4, extend CONTRIBUTING.md. Add a CUDA/Metal overlay subsection to the existing "Bumping the MLX upstream pin" section (lines 127-153 today), beside the ROCm paragraph, giving CUDA/Metal what ROCm already has: the base-declaration convention, the drift-check block below, the 3-way-merge procedure, and the rule that a bump updates every rebased overlay's base line. Do not restate the pin mechanics already documented there, and do not duplicate #1813's ROCm procedure. Fix the stale path while there: the fused Metal launchers live in src/lib/mlx-cpp/turbo/.

Scope

In scope: base lines in all 30 files under src/lib/mlx-cpp/patches/ and src/lib/mlx-cpp/patches-cuda/; the rebase of the four files above; the GIT_TAG line at src/lib/mlx-cpp/CMakeLists.txt:171; one new subsection in CONTRIBUTING.md.

Out of scope: ROCm retargeting and the ROCm fork sync script (#1813; patches-rocm/UPSTREAM still records retargeted_to_mlx_pin: 81ba1c6a, so either bump it under #1813 or note in this PR that ROCm still points at the old pin); a CI script or gate that runs the drift check automatically (follow-up, only possible once the bases exist); the mlxcelverse reorg (#1816); retiring overlays upstream has absorbed, unless the drift check shows one is now a pure revert.

Implementation Notes

  • Reuse: scripts/ci/mlx_pinned_commit.sh to read the pin (do not re-implement the parser), git merge-file as in patches-rocm/README.md, the header convention in patches/mlx/backend/cuda/cuda_utils.h, and patches-rocm/LOCAL_FIXES.md as the model for recording what changed.
  • Constraints: overlays are whole-file, so a small-looking bump can revert a large upstream change; build.rs purges _deps/ when the pin changes (_deps/.mlx-build-commit), so the first build after Phase 3 is a full MLX rebuild; the checkout the drift check needs is $OUT_DIR/build/_deps/mlx-src (src/lib/mlxcel-core/build.rs:343), or any local ml-explore/mlx clone.
  • Edge cases: an overlaid file deleted or renamed upstream (drop or move the overlay and say so); an overlay upstream now fixes natively (retire it, per chore: bump vendored MLX to upstream main (2026-06-11) and retire the steel GEMM overlay #222 and chore(deps): report the CUDA RMSNorm small-axis regression upstream and drop the local overlay once fixed #830); a tag-only base that does not resolve in a clone without tags.
  • Error handling: if a 3-way merge cannot be resolved with confidence, stop and say so in the PR rather than hand-merging a quantized kernel by guesswork; if the CUDA gate falls below the baseline, do not merge on the assumption that the failures are unrelated.

Acceptance Criteria

  • All 30 files under src/lib/mlx-cpp/patches/ and src/lib/mlx-cpp/patches-cuda/ carry exactly one Modified from upstream MLX <sha> <path> line whose sha resolves in an MLX checkout, each established by diff rather than stamped.
  • patches/mlx/backend/cuda/device/qmm_naive.cuh, patches/mlx/backend/cuda/quantized/qmm/qmv.cu, patches/mlx/backend/cuda/quantized/quantized.cpp and patches/mlx/backend/metal/quantized.cpp are rebased onto the new pin with every documented mlxcel delta intact, and the dfe17bafb (upstream #3912) fix is present in the rebased result, not reverted; the PR states how that was confirmed.
  • GIT_TAG at src/lib/mlx-cpp/CMakeLists.txt:171 is the new 40-character pin, bash scripts/ci/mlx_pinned_commit.sh prints it, and cargo test -p mlxcel-mlx-pin passes.
  • The drift check reports 0 upstream commits for all 30 overlays against the new pin, and no NO BASE line.
  • CONTRIBUTING.md documents the CUDA/Metal base convention, the drift check and the 3-way merge, without duplicating the ROCm paragraph or chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813.
  • The full CUDA gate on GB10 is at or above today's baseline on main: 11,369 passed, 0 failed, 126 suites, 0 aborts.
  • The four rebased kernels are verified against a real quantized checkpoint, with decode throughput compared before and after and reported in the PR.
  • The three Metal turbo launcher tests are either run on Apple Silicon and reported, or explicitly flagged as unrun with the reason.
  • All four phases land in a single PR, with the local gate run once at the end and pushes kept to a minimum, since any Rust-touching push costs a full CI gate cycle.

Verification

# Pin reads back through the one parser that matters
bash scripts/ci/mlx_pinned_commit.sh
cargo test -p mlxcel-mlx-pin

# Overlay drift: any output means an overlay lacks a base or may be reverting upstream work
MLX_SRC=<ml-explore/mlx clone, or $OUT_DIR/build/_deps/mlx-src>
PIN=$(bash scripts/ci/mlx_pinned_commit.sh)
for f in $(find src/lib/mlx-cpp/patches src/lib/mlx-cpp/patches-cuda -type f); do
  case "$f" in
    */patches-cuda/*) up="mlx/$(basename "$f")" ;;
    *) up="${f#src/lib/mlx-cpp/patches/}" ;;
  esac
  base=$(sed -n 's|^// *Modified from upstream MLX \([0-9a-f]\{7,40\}\).*|\1|p' "$f" | head -1)
  [ -n "$base" ] || { echo "NO BASE $f"; continue; }
  n=$(git -C "$MLX_SRC" log --oneline "$base..$PIN" -- "$up" | wc -l)
  [ "$n" -eq 0 ] || echo "DRIFT $n $f (base $base)"
done

# Full CUDA gate on GB10. --test-threads=1 is mandatory, not stylistic: the whole
# CUDA suite aborts at the default thread count (#1048, and
# src/lib/mlxcel-core/src/cuda_test_serialization_tests.rs documents it).
cargo test --workspace --profile test-fast --features cuda --no-fail-fast -- --test-threads=1

Cross-host gate, Apple Silicon only, impossible on the GB10 CUDA box: the three in-tree fused Metal kernel launcher tests for src/lib/mlx-cpp/turbo/, sparse_v_kernel_threshold_zero_matches_graph, delegated_fused_kernel_matches_reference_over_200_steps and delegated_steel_envelope_matches_cold_only_fused_over_200_steps, each within RMS < 5e-3 of the graph reference. Run them on M-series hardware, or state in the PR that they are unrun and why.

Manual, on a real quantized checkpoint (download one if absent): greedy decode output before and after the rebase for correctness, then a decode-throughput comparison. Per repo policy, take throughput numbers once at the end of the run on an otherwise quiet box, not interleaved with the test gate.

Technical Considerations

Related: #1047 (pin as single source of truth), #1042 (the previous bump-and-rederive, where 3 of 21 overlays conflicted), #1813 (ROCm fork sync and ROCm pin-bump procedure, deliberately not covered here), #1816 (mlxcelverse reorg, which may relocate patches* later), #830 and #222 (overlays retired once upstream carried the fix).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificplatform:macosmacOS (Apple Silicon) specificpriority:mediumMedium prioritystatus:readyReady to be worked ontype:choreMaintenance tasks (build, CI, etc.)

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions