Skip to content

update: Advance MLX pin to 243bdf09 and reconcile backend overlays #2303

Open

Description

@inureyes

Problem / Background

Advance MLX from 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 to upstream main captured on 2026-10-14 KST, 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e. The range contains 111 commits and changes 312 files; a pin-only edit can silently replace upstream fixes with stale whole-file overlays. #1932 / #2298 completed the previous bump, not this update.

Current Behavior

At mlxcel main 4b0c4519, src/lib/mlx-cpp/CMakeLists.txt:18 defines mlx_apply_source_overlays, called before adding MLX at line 213; line 203 is the sole executable pin. Its 32 CPU/CUDA/Metal overlays under patches/ and patches-cuda/ carry base declarations. src/lib/mlxcel-core/build.rs:296 tracks their rebuild inputs; build_support/mlx_pin.rs owns pin parsing and stale-cache checks. ROCm independently overlays MLX core files and records its target in src/lib/mlx-cpp/patches-rocm/UPSTREAM:5.

Proposed Solution

Implement all four items in one PR because they share the dependency boundary and hardware verification; run the complete local gate once after reconciliation and keep pushes to a minimum.

  1. Reconcile CPU/CUDA/Metal overlays. Follow CONTRIBUTING.md:150: inventory every overlay and documented local delta, compare its declared base with pristine target sources, three-way merge changed files, and inspect the residual against upstream. Preserve required custom kernels, dtype behavior and runtime controls. Update provenance only for files actually rebased; unchanged bases remain valid when path-specific drift is zero. Remove redundant overlays and stale CMake registrations, and handle upstream renames explicitly. Record retained/retired deltas with upstream evidence in the PR.
  2. Audit fixes now upstream. Check the B * M gathered-row correction in [nightly-verify] main is red #1599 / PR fix(metal): size gather_qmm_rhs by the broadcast x rows #2138 against target Metal quantized.cpp; check its pending Xcode 27 gated-delta NAX workaround against upstream #4594. Preserve behavior without blindly copying the older entire files. Check patches/mlx/backend/cuda/cuda_utils.h:40 teardown handling against upstream #4514 (9eb3e3ab6fceb4031994f70dfe023872643ed634), which superseded closed, unmerged upstream #4531. Retire local code only after confirming equivalent error/lifetime semantics and preserving unrelated deltas.
  3. Retarget ROCm and the pin. Set CMake's sole GIT_TAG to the captured target. Reuse scripts/mlxcelverse/rocm_overlay.py retarget, then review drift and regenerate CORE_RESIDUAL.diff; update UPSTREAM, LOCAL_FIXES.md and counts when necessary. Keep fork commit 75915908dfe5028335d318b10340313744fd3a8d fixed. Adapt new primitive signatures/fallbacks as needed; do not drop the ROCm backend or hide unresolved symbols. Preserve existing fail-closed pin/cache checks and backend feature isolation. The target changes allocator::malloc to move-only allocator::Data, makes array::set_data consume Data, checks ownership in is_donatable, and removes the free wrapper. Adapt allocation/data handoff sites in custom kernels, src/lib/mlxcel-core/cpp/, src/lib/mlx-cpp/turbo/fused_rope_append.h and the ROCm backend, including mlx/backend/rocm/allocator.{h,cpp} and mlx/backend/common/compiled.{h,cpp}; preserve asynchronous buffer lifetime, custom release behavior and donation safety rather than casting away ownership. Reconcile SDPA has_sinks, new GatedDeltaUpdateVJP and FloorDivide primitives, and FFT odd_out with supported ROCm execution or explicit graph fallbacks.
  4. Validate actual execution. Build against the new fetched HEAD, run affected numerical tests and real quantized inference on CUDA, and run candidate Metal CI on Apple Silicon. Coordinate with fix(kimi_linear): diagnose CUDA token-zero degeneration #2292 / PR fix: restore Kimi Linear CUDA inference #2296 for normal Kimi inference: main still lacks its additive-mask fix. If testing a combined validation tree, record both exact source heads and binary hash; do not attribute that model fix to this pin update.

Scope and Failure Handling

This is the dependency refresh and necessary compatibility work, not the #1816 layout refactor or an implicit closure of #1599, #2158 or #1986. Preserve open sibling PR work. Do not accept unresolved merge markers, missing upstream changes, changed numerical behavior without analysis, or unrun hardware checks as passes. Record unavailable ROCm/M5 coverage explicitly. Stop at a reviewable PR; this issue does not authorize merging.

Acceptance Criteria

  • Pin parser and actual fetched MLX HEAD equal the captured target; every overlay has reviewed provenance and no unexplained path drift or lost local functionality.
  • ROCm records and residual match the target; required backend hooks remain integrated. Allocation/data transfers compile with move-only ownership, retain buffers until asynchronous completion, release exactly once, and reject donation of borrowed storage. Add meaningful ownership/lifetime regression coverage where existing tests do not exercise the adapted path; include empty arrays and failure cleanup.
  • CUDA compilation, affected numerical regressions and real Kimi short/long greedy prompts pass with provenance; output is finite and meaningful, without token-zero repetition. Retain the fix: restore Kimi Linear CUDA inference #2296 exact short answer 4, long-context marker AMBER-RAVEN-7429, and natural-language smoke with default SDPA as controls.
  • Candidate Metal build/tests run on the new source, including models::switch_layers::mxfp_tests shared-activation cases and the three turbo launcher tests documented in CONTRIBUTING.md:142; old-pin results are not substituted.
  • PR records hardware, source/pin/binary hashes, failures and remaining unrun checks. The previously requested combined integration gate runs once after the expanded fixes are ready.

Verification

bash scripts/ci/mlx_pinned_commit.sh
cargo test -p mlxcel-mlx-pin
make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-rocm-overlay verify-python-tooling verify-fmt
cargo test --locked -p mlxcel -p mlxcel-surgery --lib --profile test-fast --features cuda --no-run
cargo test --locked -p mlxcel --lib --profile test-fast --features cuda kimi_linear -- --test-threads=1
cargo test --locked -p mlxcel --lib --profile test-fast --features cuda models::switch_layers::mxfp_tests -- --test-threads=1

Run the provenance drift loop from CONTRIBUTING.md:161 and ROCm retarget/drift procedure from docs/mlxcelverse/rocm-overlay.md. On Apple Silicon, dispatch .github/workflows/nightly-verify.yml for the candidate branch and record its exact tested commit; it runs Metal verification and video tests, but carries no model weights. On an AMD host, run scripts/mlxcelverse/check_api_drift.sh and make verify-rocm, or explicitly report unavailable hardware. At the final coordinated integration gate, run make verify-test-cuda, workspace CUDA Clippy and a fresh main Metal nightly. Serialize GPU execution with the host GPU lock and single-threaded CUDA tests.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions