Repository navigation
chore(mlx): bump pin to e9463bb and rebase CUDA overlays (#625) - #642
Merged
Merged
Conversation
Bump the pinned MLX commit from a6ec7123dac814417147e21d4aeed694924ddd4d (2026-06-10) to e9463bbfc1a7cd9e0e6b96aaa3068a316e234a63 (upstream main, 2026-07-01) in both enforcement points (mlx-cpp/CMakeLists.txt GIT_TAG and mlxcel-core/build.rs MLX_EXPECTED_COMMIT), and reconcile every source overlay three-way against the 63 upstream commits between the pins (qmv global scale #3723, JIT qmm/gather_gemm #3706 and qmm_naive #3576, rope without copy #3704, Tegra managed-memory gate #3701, fused SDPA asymmetric head dims #3637, NAX qmm fixes #3631/#3632, large-uncontiguous-grid fix 4885acd).
git diff --stat (src/lib/mlx-cpp/ overlay rebase plus supporting changes):
.github/workflows/release.yml | 40 +-
.../mlx-pin-upgrade-2026-07-03.md | 139 ++++++
docs/installation.md | 30 +-
src/bin/mlx_server.rs | 5 +
src/lib/mlx-cpp/CMakeLists.txt | 25 +-
src/lib/mlx-cpp/patches-cuda/ops.cpp | 362 +++++++++++----
.../patches/mlx/backend/cuda/binary/binary.cuh | 136 +++---
.../patches/mlx/backend/cuda/jit_module.cpp | 66 ++-
.../patches/mlx/backend/cuda/primitives.cpp | 41 --
.../patches/mlx/backend/cuda/quantized/qmm/qmm.h | 149 ------
.../patches/mlx/backend/cuda/quantized/qmm/qmv.cu | 504 ---------------------
.../mlx/backend/cuda/quantized/quantized.cpp | 25 +-
src/lib/mlxcel-core/build.rs | 3 +-
src/main.rs | 3 +
src/models/gemma4_mtp_target_tests.rs | 44 +-
src/models/sanitize.rs | 28 +-
16 files changed, 707 insertions(+), 893 deletions(-)
Per-file decisions:
- patches/mlx/backend/cuda/quantized/qmm/qmm.h: dropped (superseded), the overlay delta vs the old pin was a documentation comment only and pristine e9463bb adds the global_scale parameter from upstream #3723 which we want unmodified.
- patches/mlx/backend/cuda/quantized/qmm/qmv.cu: dropped (superseded), comment-only delta vs the old pin (the broadcast_w fix it documented had already landed upstream), and pristine e9463bb carries qmv global-scale support (#3723) plus the relocated device/cute_dequant.cuh include (#3576).
- patches/mlx/backend/cuda/primitives.cpp: dropped (no-op), the overlay was byte-identical to upstream apart from an annotation comment.
- patches/mlx/backend/cuda/quantized/quantized.cpp: rebased, re-applied the ensure_row_contiguous fix for w/scales/biases in QuantizedMatmul::eval_gpu (non-contiguous 3D batched MLA weights, upstream-rejected) onto the new base that now passes std::nullopt global_scale to qmv; the GatherQMM path returns to pristine upstream.
- patches/mlx/backend/cuda/jit_module.cpp: rebased, re-applied the /proc/self/exe executable-dir resolution, the MLXCEL_CCCL_DIR override, and the cold-JIT stderr notice onto the new base, adopting upstream's cccl_dir() fallback, the extra bundled include path, the JIT'd qmm/cute_dequant header registration (#3706/#3576), the CUTLASS nvrtc flags, and the get_jit_module(Device&) signature; MLX_PTX_CACHE_DIR handling is unchanged upstream so the persistent PTX cache redirect still takes effect (upstream #3587 turned out to be CI-cache-only, no code conflict).
- patches/mlx/backend/cuda/binary/binary.cuh: rebased (adapted), re-applied the mixed-precision bf16/fp32 kernels, trait, and dispatch onto the new base; this also drops stale v0.31.1-era reversions the old overlay was silently carrying (missing __launch_bounds__/BINARY_MAX_BLOCK_DIM) and adopts upstream's large-grid index_rest fix (grid.z) and get_launch_args_general in both the same-type and the mixed general paths.
- patches-cuda/ops.cpp: rebased (adapted), re-applied bf16_mixed_astype() and its use in add/subtract/multiply/divide/maximum/minimum onto the new base; this drops stale v0.31.1-era reversions the old overlay was silently carrying (safe_cast overflow checks, the arange zero-step check, use-after-move-prone auto& shape bindings) and picks up upstream's new array-API ops unmodified.
- src/lib/mlx-cpp/CMakeLists.txt: pin bump, retargeting the MLX_CCCL_DIR compile definition to the new mlx_dirs OBJECT target (upstream moved the macro consumer into dirs.cpp) with a fallback to the mlx target, and a new MLX_CUTLASS_DIR build-tree fallback because the JIT'd quantized kernels now need CUTLASS/CuTe headers at runtime.
- Runtime JIT deployment: upstream #3706/#3576 moved the CUTLASS qmm/gather_gemm instantiations from build-time nvcc to runtime NVRTC (cold build dropped from 30+ min to ~7.5 min on GB10), so quantized models now need cute/ and cutlass/ headers on the deployment host at first run; the jit_module overlay resolves them via MLXCEL_CUTLASS_DIR, the bundled include/ layout, or the compiled-in build-tree path, release archives bundle include/cute + include/cutlass next to include/cccl, and docs/installation.md documents the requirement.
- src/lib/mlxcel-core/build.rs: also watch ../mlx-cpp/patches-cuda with rerun-if-changed (it was unwatched, so edits to the CUDA-only dtype/fast/ops overlays did not retrigger the MLX build).
- Unchanged overlays (base file identical between the two pins, kept as-is): device/binary_ops.cuh, gemms/grouped_gemm.h, gemms/grouped_gemm_unaligned.cu, matmul.cpp (CUTLASS GatherMM rework intentionally still bypasses upstream's JIT gather_gemm; re-baseline sweep will confirm), reduce/{all,col,init,row}_reduce.cu, reduce/reduce_ops.cuh, metal/compiled.cpp, metal/kernels/utils.h, patches-cuda/{dtype,fast}.cpp.
Test-suite reconciliation (cargo test --release --features cuda on GB10):
- src/models/sanitize.rs: fix NVFP4 block-scale parsing to be load-dtype-agnostic (astype to F32 before the raw-byte parse). The CUDA loader widens F16 safetensors to F32 at load (from_bytes_f16, CUDA branch since 2026-03), so the previous F16 byte parse read zeros on CUDA; this is a pre-existing CUDA bug exposed by running the suite here, not a pin regression, and it also corrects real gemma4-nvfp4 scale handling on CUDA. f16_to_f32 is now cfg(test) (kept for its format unit tests).
- src/models/gemma4_mtp_target_tests.rs: the batched-vs-B=1 bitwise-equality assertion now uses a 2e-3 fp32 tolerance on CUDA (bitwise retained elsewhere). Since MLX #3704 (rope without copy) RoPE preserves input layouts instead of canonicalizing with a copy, so downstream GEMM/SDPA kernels can see different strides for the batched row than for the B=1 replay and pick different reduction orders (observed max_abs 1.6e-4 on fp32 hidden of magnitude ~8). The argmax-equality assertion is unchanged, preserving the functional MTP verify guarantee.
- docs/installation.md: document running the CUDA test suite with --test-threads=1. The JIT'd kernels are now compiled and module-loaded on first use, and those first-use paths are not safe against concurrent test threads: the default parallel run aborts with cudaStreamEndCapture capture errors (module load racing another thread's capture) or, with MLX_USE_CUDA_GRAPHS=0, with cuLaunchKernelEx invalid-argument (kernel-configure race). Inference binaries are unaffected.
- src/main.rs: add the GLM VLM, Kimi VLM, and PaddleOCR VLM families to FAMILY_ORDER (pre-existing gap on main from the recent VLM ports; family_order_is_exhaustive fails without it and blocks the suite).
- src/bin/mlx_server.rs: the legacy-server-mode resolver test fixture now writes a stub model.safetensors next to config.json; the offline completeness gate from PR #604 classifies a config-only directory as an interrupted download and re-fetches from the network, which 404s on the synthetic repo id (pre-existing on main since #604).
Tegra note: upstream #3701 gates managed memory on concurrentManagedAccess only; GB10 reports concurrentManagedAccess=1, so managed memory stays enabled on GB10 under both pins and #3701 is a no-op on this platform (details in docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md).
Adds the focused GB10 re-baseline CSV (benchmarks/cuda_gb10_2026-07-03.csv, 19 models covering all Phase-3 outlier families plus dense/MoE/SSM/VLM representatives and a pure-mamba control) and fills the outlier-verdict table in docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md. Verdicts: MoE prefill collapse (#629) is unchanged (our matmul.cpp CUTLASS GatherMM overlay bypasses upstream's now-JIT'd gather_gemm, so #3706 does not reach the path); the hybrid-SSM decode gap (#631) is unchanged (pure-mamba control healthy, isolating the gap to the hybrid path); the gemma-4-31b nvfp4 outlier (#630) is root-caused (the checkpoint is NVIDIA ModelOpt-packed and now correctly rejected; the genuine MLX-native nvfp4 dequant bug was fixed here in sanitize.rs; mxfp4 sibling healthy) and re-scoped to end-to-end validation plus a guard test. No Phase-3 issue is fully fixed by the pin bump, so Phase 3 stands at three issues. Dense decode is uniformly +1 to +3% with several prefill gains and no cross-family regression.
This was referenced Jul 2, 2026
15 tasks
11 tasks done
inureyes
added a commit
that referenced
this pull request
Aug 6, 2026
## Summary `cargo test --release --features cuda -p mlxcel-core --lib` has never reached a verdict on GB10. It dies with SIGABRT partway through, at a different test and with a different CUDA error each run, so the abort reads as if whichever test happened to be running is broken. This PR makes the CUDA suite reach a verdict, makes serialization the property of the command rather than of the developer's memory, and adds the CUDA test gate that did not exist at all. ## What the measurements say Three full runs of that command on `main` at `6276f556`, release profile, `--features cuda`, idle GB10 (sm_121), MLX pin `2c46b953`: | Run | Threads | `MLX_USE_CUDA_GRAPHS` | Outcome | |---|---|---|---| | 1 | default (20) | on | SIGABRT, `cudaStreamEndCapture ... failed: operation failed due to a previous error during capture` | | 2 | `--test-threads=1` | on | **ran to a verdict: 1406 passed, 4 failed, 88.02s** | | 3 | default (20) | `0` | SIGABRT, `cuLaunchKernelEx ... failed: invalid argument` | **Run 3 is the result the issue did not anticipate, and it is the one that decides the fix.** Disabling graph capture does not stop the abort. It only changes which CUDA call reports it: under capture the poisoned state surfaces at `cudaStreamEndCapture`, without capture the same race surfaces as an invalid `cuLaunchKernelEx`. So the cause is the libtest harness driving MLX concurrently from many host threads, and graph capture merely selects the symptom. The issue floated "disable capture for the test binary, for example by setting `MLX_USE_CUDA_GRAPHS=0` for the suite" as a candidate remedy. It is measurably not one. This PR says so in the guard's panic message, the Makefile comment, CONTRIBUTING.md and docs/installation.md, because a future reader hitting this abort will otherwise reach for exactly that. Two consequences worth stating plainly: - **The acceptance criterion "at least one test still exercises CUDA graph capture if capture is disabled for the rest of the suite" is moot, not skipped.** Nothing gets disabled. Run 2 has capture fully on and all 1410 tests run under it, so the whole suite keeps exercising capture. There is no carve-out to arrange. - **Machine state is not the explanation**, which the issue kept open as an alternative. Runs 1 and 3 abort on an idle, non-benchmarked machine and run 2 passes on the same machine minutes apart. The outcome tracks the thread count, not the machine's history. Serialization is cheap: 88 seconds for 1410 tests. There is no throughput argument against making it the gate. Mechanism, consistent with the note already in docs/installation.md from #642: MLX begins capture with `cudaStreamCaptureModeThreadLocal` (`mlx/backend/cuda/device.cpp` at the pinned commit), which restrains only the capturing thread and leaves the others free to issue capture-unsafe calls, while the allocator issues `cudaMalloc`/`cudaFree` on demand from whichever thread needs memory. That accounts for the moving abort site, the varying message, and scoped runs passing on the same binary. ## The second hole: there was no CUDA gate to fix `make test-fast-cuda` already passed `--test-threads=1`, but it omits `--workspace`. Per the workspace comment in the root `Cargo.toml`, the workspace root is itself the `mlxcel` package, so a bare `cargo test` resolves to `-p mlxcel` and never builds `mlxcel-core`. That target cannot run a single one of mlxcel-core's 1410 tests. `make verify-test` is `--workspace` but pins `--features metal,accelerate`, so it is the macOS gate. Net: on Linux/CUDA no target ran mlxcel-core's tests, in the crate that holds the MLX `cxx` bridge, `layers.rs`, the KV cache and the quantization loaders. That is the blindness #1007 removed on macOS, still present on the other backend, and CUDA is gated nowhere in CI, not even nightly, so nothing else was covering it. ## What changed **`make verify-test-cuda`** is the Linux/NVIDIA counterpart of `make verify-test`, following the shape #1007 established so the two gates stay symmetrical: `cargo test --workspace --profile test-fast --features cuda --no-fail-fast -- --test-threads=1`. Same scope, same profile, same `--no-fail-fast` reporting, `cuda` in place of `metal,accelerate`, and the serialization flag embedded so the documented command cannot be run wrong. On the `--workspace` question the issue raised: pulling in `mlxcel-surgery` and `mlxcel-xla` on the CUDA path **is** what we want, and it is close to free. Both depend on `mlxcel-core`, and one `cargo test --workspace --features cuda` unifies that into a single `cuda`-enabled build of it, so neither triggers a second MLX compile. `mlxcel-xla`'s `iree` feature stays off, so it remains pure Rust and needs no IREE distribution. `mlxcel-mlx-pin` does not depend on `mlxcel-core` at all. Their test targets are backend-agnostic Rust, so excluding them would only mean those two crates are gated on macOS and nowhere else. This is written down in docs/installation.md next to the existing macOS explanation. There is deliberately **no** `verify-clippy-cuda` in this PR. The CUDA lint hole is a separate hole from the CUDA test hole, #1048 is about the test gate, and `cargo clippy --workspace --all-targets` does not complete in a usable time on this host, so I would have been adding a target I could not run. **A harness-level guard**, `the_cuda_test_suite_must_run_single_threaded` in the new `src/lib/mlxcel-core/src/cuda_test_serialization_tests.rs`, is the part that does not depend on remembering a flag. It resolves the thread count the way libtest itself does (command line `--test-threads` wins over `RUST_TEST_THREADS`, otherwise `available_parallelism()`) and fails by name, with the right command, whenever the CUDA suite is about to run multi-threaded. A hand-run `cargo test --workspace --features cuda` now gets a diagnosis instead of an anonymous abort. Three properties worth noting: - Being an ordinary test, it is filtered out of narrowed runs, which stay parallel. That is correct rather than a loophole: scoped subsets pass parallel, and it is whole-suite runs that abort. - It writes the message straight to the stderr handle rather than through `eprintln!`. libtest captures the print macros and replays them in its end-of-run summary, and in a parallel run there may be no end of run, because the abort takes the process down first. Writing to the handle bypasses the capture. Verified: the message appears on a 20-thread run. - `MLXCEL_ALLOW_PARALLEL_CUDA_TESTS=1` downgrades it to a printed warning, mirroring `MLXCEL_ALLOW_CONCURRENT_GPU_TESTS` in `gpu_exclusivity_tests.rs`, for someone deliberately reproducing the abort. It is gated on `feature = "cuda"`, so macOS is untouched. The Metal hazard is the cross-process one in #1008, which `gpu_exclusivity_tests` already covers, and `--test-threads` capping is documented there as not helping it. **Documentation.** CONTRIBUTING.md gains the CUDA gate command, the requirement that it be serialized, and the run-3 reason not to reach for `MLX_USE_CUDA_GRAPHS=0`; its Linux/CUDA sample command also gains the `--test-threads=1` it was missing, so following the contributor docs literally no longer produces a SIGABRT. docs/installation.md gains the measured table, the negative result, the guard's location and the CUDA feature-resolution note. Both state that `test-fast` / `test-fast-cuda` are edit-test-loop targets and not gates. docs/environment-variables.md gains `MLXCEL_ALLOW_PARALLEL_CUDA_TESTS` in the Test and CI table, along with `MLXCEL_ALLOW_CONCURRENT_GPU_TESTS`, which the #1008 guard has always had and which was never documented there either; that file is where people grep for what an `MLXCEL_*` variable does, and an existing omission is a weak argument for adding a second. ## Why not the other options - **`.cargo/config.toml` `[env] RUST_TEST_THREADS = "1"`.** Cargo's `[env]` cannot be conditioned on platform or feature, so this would also serialize the macOS gate, where serializing buys nothing because #1008's hazard is between processes rather than between threads, and I cannot measure that cost from this host. It would also move the gate's behavior into a file the gate's own command line does not mention, which is precisely the objection the root `Cargo.toml` already records against using `default-members` for scope. - **A committed nextest profile.** nextest runs each test in its own process, which recreates #1008's cross-process condition, the worse of the two failure modes, and it adds a required tool to the gate. - **The Makefile target and documentation alone.** That covers the documented command but leaves a hand-run `cargo test --workspace --features cuda` aborting anonymously, which is the case the guard exists for. ## Known-failing tests: the gate is not green once it runs Run 2 reports 4 failures. None are fixed here, and none are marked `#[ignore]`, because hiding them would recreate exactly the blindness this issue exists to remove. | Test | Location | Message | Tracked | |---|---|---|---| | `fused_rope_parity_tests::fused_rope_append_matches_graph_rope_across_offsets` | `src/lib/mlxcel-core/src/fused_rope_parity_tests.rs:87` | `q offset=131071: normalized rms 2.684e-3 (tol 2.0e-3), normalized max 3.663e-2 (tol 1.2e-2)` | #1049 | | `fused_rope_parity_tests::fused_rope_append_multi_token_positions_are_absolute` | `src/lib/mlxcel-core/src/fused_rope_parity_tests.rs:87` | `q seq=3 offset=131069: normalized rms 2.497e-3 (tol 2.0e-3), normalized max 3.698e-2 (tol 1.2e-2)` | #1049 | | `paged_v2::launch::launch_tests::chunk_size_does_not_change_the_answer` | `src/lib/mlxcel-core/src/paged_v2/launch_tests.rs:324` | `pages_per_chunk 1 relative error 55896.2` | #1054 | | `paged_v2::sparse::sparse_tests::an_f16_cache_matches_the_reference_within_its_own_precision` | `src/lib/mlxcel-core/src/paged_v2/sparse_tests.rs:424` | `max relative error 1.0004048 on an f16 cache` | #1053 | Every failure a red CUDA gate can currently report maps to a tracking issue, so the next person reading one does not have to work out which of them are accounted for. The two `paged_v2` failures were unfiled when this branch was opened and are now #1053 and #1054, split rather than filed together because the triage below says they are different bugs. From a scoped serialized run on this branch (`-p mlxcel-core --lib paged_v2 -- --test-threads=1`, 91 passed, 1 failed): - `an_f16_cache_matches_the_reference_within_its_own_precision` (#1053) **reproduces in isolation**. A relative error of essentially 1.0 means the output carries no relation to the reference, so this is a correctness failure rather than tolerance drift, against a tolerance already sized for f16 rounding. The `assert!(outcome.is_fused(), ...)` at `sparse_tests.rs:421` runs before the numeric assert at 424, so the fused sparse kernel path is genuinely taken and the failure is purely numeric. - `chunk_size_does_not_change_the_answer` (#1054) **passes** both inside that 92-test scope and run entirely alone (1.73s). It fails only in the full suite. That points at leaked shared mutable state rather than a defect in the chunking path itself, which is a materially different investigation from the first one. The assert at `launch_tests.rs:324` sits inside `for ppc in [1, 2, 3, 4, 7, 16, 64]`, so the reported `pages_per_chunk 1` is only the first iteration and the larger chunk sizes are never reached; the failure is not specific to chunk size 1. Determining whether either also fails on a CPU-only build would need a separate non-CUDA MLX build, which I did not run. ## Test plan The gate itself, on GB10, run end to end. - [x] `make verify-test-cuda`: the whole workspace under `[profile.test-fast]` with `--features cuda --no-fail-fast -- --test-threads=1`, cold-building that profile first. `mlxcel-core` lib: **1411 passed, 4 failed, 89.35s, no SIGABRT.** Every other workspace target passed, roughly 6674 further tests across the other binaries. `make` exited 2 on the four failures, which is what a red gate should do. Everything below on the same host, release profile, `--features cuda`, narrow targets only. - [x] Guard passes serialized: `cargo test --release --features cuda -p mlxcel-core --lib cuda_test_serialization -- --test-threads=1` gives 5 passed, 0 failed. - [x] Guard fails as intended at 20 threads: same command without the flag gives 4 passed, 1 failed, with the full diagnosis and both commands to run. - [x] The diagnosis survives libtest's capture: the direct stderr write appears in the 20-thread run's output. - [x] Escape hatch works: `MLXCEL_ALLOW_PARALLEL_CUDA_TESTS=1` on the 20-thread run gives 5 passed with a printed warning. - [x] Scoped runs are unaffected: `-p mlxcel-core --lib paged_v2 -- --test-threads=1` runs 92 tests with the guard filtered out. - [x] `cargo fmt --all -- --check` clean. - [x] `cargo clippy -p mlxcel-core --release --tests --features cuda -- -D warnings` clean, so the new module's test target carries no lint debt (this is the target `--all-targets` without `--workspace` would not have compiled). - [x] `cargo tree --workspace --features cuda` resolves all five members, so the gate's feature selection is valid. - [x] `make -n verify-test-cuda` emits the intended command; `make help` lists the new target. - [x] `python3 scripts/ci/check_cross_repo_refs.py` clean. ## What the end-to-end gate run settles Two gaps this PR opened with, both now closed by that run. **The serialization result carries across profiles.** The three diagnostic runs were taken under `[profile.release]`, while the gate builds under `[profile.test-fast]`. Those differ in LTO and codegen units, and this repository has a standing note that a defect can reproduce under one and not the other, so the gate could have been resting on a profile nothing had tested. It is not: 1411 passed, 4 failed, 89.35s, no abort, against 1406/4/88.02s under release. **The known-failing list is accurate for the command a contributor actually runs**, not only for the release-profile command the diagnosis used. The same four failures reproduce under `test-fast`, at the same assertion sites: `fused_rope_parity_tests.rs:87` twice (#1049), `paged_v2/launch_tests.rs:324` (#1054), `paged_v2/sparse_tests.rs:424` (#1053). Nothing else in the workspace is red. One incidental result worth recording: no other test binary aborted either, so the root `mlxcel` suite having no guard of its own did not bite on this run. The guard covers `mlxcel-core` because that is where the aborts were measured; whether the root suite needs one is unresolved rather than answered. ## Not verified **The fresh-boot half of the acceptance criterion is still not met.** The criterion asks for a verdict "both on a fresh boot and after a benchmarking session". Every run here, the three diagnostic ones and the gate run, is on a machine that has been up for days with builds and several full suite runs the same day, which covers the "after a session" half and arguably covers it harder than the criterion intended. A fresh-boot run needs a reboot, which is the maintainer's call, and prior sessions on this host recorded driver-level instability accumulating after heavy runs. Closes #1048
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
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
Bumps the pinned MLX commit from
a6ec7123dac814417147e21d4aeed694924ddd4d(2026-06-10) toe9463bbfc1a7cd9e0e6b96aaa3068a316e234a63(upstream main, 2026-07-01) in both enforcement points, and reconciles all 20 source overlays three-way against the 63 upstream commits between the pins. Part of epic #623; the full GB10 re-baseline sweep and outlier triage verdicts are done separately on top of this (placeholder committed atdocs/benchmark_results/mlx-pin-upgrade-2026-07-03.md).Pin bump
src/lib/mlx-cpp/CMakeLists.txtGIT_TAG->e9463bb...src/lib/mlxcel-core/build.rsMLX_EXPECTED_COMMIT->e9463bb...(the mismatch check and.mlx-build-commitpurge forced the clean MLX rebuild)MLX_CCCL_DIRis now defined on the newmlx_dirsOBJECT target (upstream moved the macro consumer intodirs.cppto isolate dynamic defines from the compile cache), with a fallback to themlxtarget for older trees. Verified in the generatedmlx_dirs.dir/flags.make.Overlay reconciliation
14 of 20 overlay base files are byte-identical between the two pins, so those overlays apply unchanged. The remaining six were reconciled as follows:
quantized/qmm/qmm.hglobal_scaleparam (#3723).quantized/qmm/qmv.cudevice/cute_dequant.cuhinclude (#3576).primitives.cppquantized/quantized.cppensure_row_contiguousfix forw/scales/biasesinQuantizedMatmul::eval_gpu(3D batched MLA weights, upstream-rejected) onto the new base, which now passesstd::nulloptglobal_scale to qmv. GatherQMM path returns to pristine upstream.jit_module.cpp/proc/self/exedir resolution +MLXCEL_CCCL_DIRoverride + cold-JIT notice onto the new base; adopted upstream'scccl_dir()fallback, extra bundled include path, JIT'd qmm header registration (#3706/#3576), CUTLASS nvrtc flags, andget_jit_module(Device&)signature.MLX_PTX_CACHE_DIRhandling is unchanged upstream, so the persistent PTX cache redirect still takes effect (upstream #3587 turned out to be a CI-cache change only).binary/binary.cuh__launch_bounds__), and adopts upstream's large-gridindex_restfix (grid.z) andget_launch_args_generalin both same-type and mixed general paths.patches-cuda/ops.cppbf16_mixed_astype()+ its use in add/subtract/multiply/divide/maximum/minimum. Drops stale v0.31.1-era reversions (safe_castoverflow checks,arangezero-step check, use-after-move-proneauto& shapebindings) and picks up upstream's new array-API ops unmodified.Kept unchanged (base identical between pins):
device/binary_ops.cuh,gemms/grouped_gemm.h,gemms/grouped_gemm_unaligned.cu,matmul.cpp(the CUTLASS GatherMM rework intentionally still bypasses upstream's JIT gather_gemm; the re-baseline sweep should confirm this remains the right call on MoE models),reduce/*,metal/compiled.cpp,metal/kernels/utils.h,patches-cuda/{dtype,fast}.cpp.Quantized kernels moved to runtime JIT (deployment impact)
Upstream #3706/#3576 JIT-compile the CUTLASS qmm/gather_gemm kernels with NVRTC at runtime instead of nvcc at build time. Cold build time dropped from the historical 30+ minutes to ~7.5 minutes on GB10, but the JIT'd kernels include
<cute/...>/<cutlass/...>, which a statically linked mlxcel binary could not resolve (first quantized run failed withcannot open source file "cute/numeric/numeric_types.hpp"). Fixed in this PR, mirroring the existing CCCL solution:jit_module.cppoverlay resolves CUTLASS headers viaMLXCEL_CUTLASS_DIRenv override, the bundled<exe>/../includelayout, or the compiled-inMLX_CUTLASS_DIRbuild-tree fallback.mlx-cpp/CMakeLists.txtcompiles in the_deps/cutlass-src/includefallback for source builds.release.ymlbundlesinclude/cute/+include/cutlass/in both Linux CUDA archive jobs (with completeness checks).docs/installation.mddocuments the new runtime requirement and env var.build.rsnow watches../mlx-cpp/patches-cudawith rerun-if-changed (previously unwatched, so edits to the CUDA-only overlays did not retrigger the MLX build).First quantized-model runs pay a one-time per-kernel-variant NVRTC compile, amortized by the persistent commit-scoped PTX cache.
Test-suite reconciliation
cargo test --release --features cudaon GB10 surfaced two failures, both understood:sanitize_tests::load_and_sanitize_weights_dequantizes_nvfp4_gemma4_checkpoint: pre-existing CUDA bug, not a pin regression. The CUDA loader widens F16 safetensors to F32 at load (from_bytes_f16, CUDA branch dating to 2026-03), while the NVFP4 dequant parsed the block scales as raw F16 bytes, reading zeros on CUDA. Fixed by normalizing the scales to F32 viaastypebefore the byte parse (dtype-agnostic on both backends). This also corrects real gemma4-nvfp4 block-scale handling on CUDA, relevant to the epic's gemma-4-31b nvfp4 outlier.gemma4_mtp_target::tests::divergent_round_hidden_matches_b1_replay: pin-related property change. Since MLX #3704 (rope without copy) RoPE preserves input layouts instead of canonicalizing them with a copy, so downstream GEMM/SDPA kernels can see different strides for a batched row than for the B=1 replay and pick different reduction orders. Observed max_abs 1.6e-4 on fp32 hidden of magnitude ~8, i.e. reduction-order noise, not a logic bug. The bitwise assertion is now a 2e-3 fp32 tolerance on CUDA (bitwise retained on other backends); the argmax-equality assertion, which is the functional MTP-verify guarantee, is unchanged.The default parallel
cargo testrun aborts under the new pin even on warm PTX caches: the JIT'd kernels are compiled and module-loaded on first use, and those first-use paths are not safe against concurrent test threads. Observed ascudaStreamEndCapture ... previous error during capture(module load racing another thread's stream capture) and, withMLX_USE_CUDA_GRAPHS=0, ascuLaunchKernelEx ... invalid argument(kernel-configure race inget_kernel_and_dims). The suite is stable with--test-threads=1; documented indocs/installation.md. Inference binaries are unaffected (parity runs and benches are clean). Worth an upstream issue; also worth watching if mlxcel-server ever evals from multiple threads.One more pre-existing failure fixed en route:
family_order_is_exhaustive(bin target) was failing on main because the recent VLM ports added the GLM VLM, Kimi VLM, and PaddleOCR VLM families without updatingFAMILY_ORDERinsrc/main.rs; the three families are now listed beforeOther VLM.Tegra managed-memory finding (#3701)
GB10 reports
concurrentManagedAccess=1(probed with a CUDA attribute query), so upstream's new gate does not trip: managed memory stays enabled on GB10 under both pins, exactly as before. #3701 is a no-op on this platform (it targets Jetson Orin Nano-class boards and WSL, which report 0). The epic hypothesis that the old pin's managed-memory behavior was depressing GB10 numbers is refuted at the gate level.Test plan
cargo build --release --features cudafrom purged_deps(fetched e9463bb, all overlays applied; verified in the configured tree andmlx_dirsflags)cargo test --features cudalocallyCloses #625