Skip to content

chore: bump vendored MLX to upstream main (2026-06-11) and retire the steel GEMM overlay #222

Closed

Description

@inureyes

Summary

The vendored MLX is pinned at 84961223c02925bef6bef95d3a0a046779bde935 (upstream main, 2026-05-09). During issue #217 the real-model gate root-caused output corruption to an upstream bug in that pin (steel GEMM BaseMMAFrag<T, 8, 8>::load_safe indexes columns with off_x instead of off_y; fixed upstream 2026-05-18/19 in ml-explore/mlx PRs #3560 and #3565), and we now carry the fix as a whole-file overlay at src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/steel/gemm/mma.h. That overlay is deliberate technical debt: once the pin moves past 2026-05-19 it MUST be deleted, because the CMake overlay mechanism (mlx_apply_source_overlays) copies our snapshot over the fetched source and would silently revert any newer upstream changes to that file.

Bump the pin to upstream main a6ec7123dac814417147e21d4aeed694924ddd4d (2026-06-11) and update everything that depends on the pinned commit.

Why now

Beyond retiring the overlay, the 2026-05-09 to 2026-06-11 window contains upstream fixes we currently lack: qvm_split_k batch stride (#3497), rope single-token multiple-sequences (#3498, the same family mlxcel worked around in #167), int32 shape-product overflow detection at compute-shape boundaries (#3524), int32 overflow in matvec row offset and gather-MM batch stride (#3609), deployment target passed when linking the metallib (#3501), steel GEMM safe-load fixes (#3560/#3565, currently overlaid), compile cache and tracing-state thread-safety fixes (#3628/#3638), and scatter_prod GPU hang on NaN (#3492).

Tasks

  • Update GIT_TAG in src/lib/mlx-cpp/CMakeLists.txt and MLX_EXPECTED_COMMIT in src/lib/mlxcel-core/build.rs to a6ec7123dac814417147e21d4aeed694924ddd4d (the two must match; build.rs purges stale _deps caches via the marker file).
  • DELETE the mma.h overlay file and its _metal_patch_files entry plus the explanatory comment in src/lib/mlx-cpp/CMakeLists.txt (the fix is upstream at the new pin).
  • Re-derive the remaining overlays against the new pin: patches/mlx/backend/metal/compiled.cpp (mlxcel delta: static_cast type promotion for mixed-dtype compiled ops) and patches/mlx/backend/metal/kernels/utils.h (mlxcel delta: #include <metal_simdgroup_matrix> and metal::vec qualification). Take the pristine upstream file at the new commit and re-apply the deltas; drop any delta upstream has absorbed.
  • Fix compile breakage from one month of upstream API drift in src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp and the turbo kernels (src/lib/mlx-cpp/turbo/, which reference MLX internals and must match the pinned commit per the build.rs warning), plus the CUDA patch files if their upstream counterparts moved.
  • Verify the kernel regression tests that were added precisely for this class of change: steel_gemm_edge_tile_safe_load_matches_reference (mlxcel-core) and the buffer-cache-warmed real_model_forward_determinism (DiffusionGemma, both forwards byte-deterministic).
  • Real-model revalidation battery: DiffusionGemma temperature-0 deterministic CLI run (expected to stay byte-stable or change ONLY via legitimate upstream numerics, to be characterized), qwen/llama paged scheduler parity suites, a Gemma 4 MTP byte-identical spot check, full serial lib suite (--test-threads=1), cold clippy on both CI feature sets.
  • Throughput sanity vs current main on M1 Ultra: qwen3-0.6b decode, gemma-4-26b-a4b decode, DiffusionGemma per-denoising-step time (651 ms per 256-canvas step baseline); flag any regression beyond noise.

Acceptance criteria

  • Vendored MLX at a6ec712, no mma.h overlay, remaining overlays re-derived and documented.
  • All listed regression and parity gates green; throughput within noise of the pre-bump baseline.
  • docs/comments that reference the old pin or the overlay updated.

Activity

  1. added
    type:enhancementNew features, capabilities, or significant additions
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    and removed on Jun 11, 2026
  2. added this to the 0.2 milestone on Jun 21, 2026
  3. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority:highHigh prioritystatus:doneCompletedtype:enhancementNew features, capabilities, or significant additions

Type

No type

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions