Repository navigation
update: Advance MLX pin to 243bdf09 and reconcile backend overlays #2303
Description
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.
- 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. - Audit fixes now upstream. Check the
B * Mgathered-row correction in [nightly-verify] main is red #1599 / PR fix(metal): size gather_qmm_rhs by the broadcast x rows #2138 against target Metalquantized.cpp; check its pending Xcode 27 gated-delta NAX workaround against upstream #4594. Preserve behavior without blindly copying the older entire files. Checkpatches/mlx/backend/cuda/cuda_utils.h:40teardown 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. - Retarget ROCm and the pin. Set CMake's sole
GIT_TAGto the captured target. Reusescripts/mlxcelverse/rocm_overlay.py retarget, then reviewdriftand regenerateCORE_RESIDUAL.diff; updateUPSTREAM,LOCAL_FIXES.mdand counts when necessary. Keep fork commit75915908dfe5028335d318b10340313744fd3a8dfixed. 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 changesallocator::mallocto move-onlyallocator::Data, makesarray::set_dataconsumeData, checks ownership inis_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.hand the ROCm backend, includingmlx/backend/rocm/allocator.{h,cpp}andmlx/backend/common/compiled.{h,cpp}; preserve asynchronous buffer lifetime, custom release behavior and donation safety rather than casting away ownership. Reconcile SDPAhas_sinks, newGatedDeltaUpdateVJPandFloorDivideprimitives, and FFTodd_outwith supported ROCm execution or explicit graph fallbacks. - 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 markerAMBER-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_testsshared-activation cases and the three turbo launcher tests documented inCONTRIBUTING.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=1Run 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.
Activity
fix(cuda): pass device to gathered GEMM tile selection
test(cuda): cover aligned sorted MXFP gathered GEMM
update(cuda): retire the upstream-covered epilogue barrier overlay