[davidtai]

Proposed changes

This PR merges upstream ml-explore/mlx main at c6abf0dbb (2026-10-22) into the fork. It uses a merge commit, not a rebase. David ordered the merge on 2026-10-22.

What upstream brings

• 272 upstream commits, from the old fork base 77a0c1e8 (2026-08-14) to c6abf0dbb.
• Tags in the range: v0.32.2, v0.32.3 and v0.32.4. The new base is 15 commits past v0.32.4. mlx/version.h is upstream's.
• 428 upstream files change (+28025 / -4806 lines).
• Metal: NAX qmv_mma / qmv_nax kernels, staged affine_qmm_t_nax with split-K, the fp4-to-half loader in fp_qmm_t, global scale for qmm and gather_qmm, NAX for non-transposed gather_qmm_rhs, 1-bit affine quantization, gated delta net kernels and VJP, SDPA VJP, SDPA D96/V64, D256 and D512 memory work, the M1/M2 tk1 SDPA kernels, fused fast.cross_entropy, the fence deadlock fixes ([BUG][Metal] Deadlock in fence ml-explore/mlx#4552, [BUG] Metal fence deadlock ml-explore/mlx#4689), conv bounds fixes, a residency refresh interval.
• Core: a FloorDivide primitive (Fix floor_divide for integers ml-explore/mlx#4515, Add FloorDivide primitive for integer floor division ml-explore/mlx#4642), stft / istft, compile fixes (Fix input/output scalar merging in compile ml-explore/mlx#4658, Fix cross-process race in CPU compile cache ml-explore/mlx#4638), complex VJP fixes, reduction grid overflow fix.
• The row-tile scheduling of gather_mm / gather_qmm ([Metal] Gather mm improvement ml-explore/mlx#4567, [Metal] gather_qmm improvement  ml-explore/mlx#4572), which the fork carried as a backport in PR [Metal] gather_mm and gather_qmm row-tile scheduling (backport of ml-explore/mlx #4567, #4572) #27.
• CUDA, Windows and distributed changes. No fork workflow builds them.

Fork changes that stay

Every fork.yaml section that describes fork work stays:

• Metal eval use-after-free (PR fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval #5)
• Sorted expert-tile quantized matmul for Gemma 4 / Qwen 3.5-3.6, with the retract to the unsorted route (PR [Metal] gather_mm and gather_qmm row-tile scheduling (backport of ml-explore/mlx #4567, #4572) #27)
• Thread-scoped command-buffer commit limits (PR [Metal] thread-scoped command-buffer commit limits #36)
• GPT-OSS MXFP4 decode masked tail and prefill tile (PR perf: accelerate GPT-OSS MXFP4 decode on M4 Max #14, PR fix(metal): build the GPT-OSS MXFP4 prefill tile into the metallib #24)
• Coherent allocator accounting and the process allocation cap (PR Expose coherent allocator accounting and correct quantized bias sums #15, PR Single-Mac GLM-5.3-Flash native fixes without JACCL #37)
• FP32 partial sums in two-pass vector SDPA (PR Preserve FP32 partial sums in two-pass vector SDPA #16)
• Declared-mutable inputs for Metal custom kernels (PR Track declared mutable inputs in Metal custom kernels #18)
• Cross-thread compile-cache cleanup (PR chore: upgrade core to MLX 0.32.2 #10)
• Validated custom Metal threadgroups (PR Single-Mac GLM-5.3-Flash native fixes without JACCL #37)
• Quantized tile bounds and distinct loader symbols (PR Single-Mac GLM-5.3-Flash native fixes without JACCL #37)
• CPU quantized accumulation precision (PR Single-Mac GLM-5.3-Flash native fixes without JACCL #37)
• The revert of the Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX ml-explore/mlx#4352 simdgroup skip in fp_gather_qmm_rhs_nax and affine_gather_qmm_rhs_nax (PR Single-Mac GLM-5.3-Flash native fixes without JACCL #37)
• Fork tooling and CI (PR feat(forkdiff): fork-diff page on GitHub Pages + CI gate that a rebase must update it #19, docs: state the upstream sync policy in FORKDIFF.md #21 to ci: remove the upstream workflows that cannot run in this fork #23, ci: report C++ line coverage from the Metal test run (re-land of #25) #28 to ci: count the Python bindings in python/src in the C++ coverage report #31)

Where the fork now takes upstream

Fork change	Upstream change	Result
PR #10, sync to v0.32.2	Upstream has v0.32.2	Section removed. Every file that PR #10 alone changed is now upstream's file.
PR #27, backport of ml-explore#4567 and ml-explore#4572	Upstream has both	Section removed. matmul.cpp, matmul.h, steel_gemm_gather*.h, gather_mm_offsets.metal, the benchmarks and .gitignore are upstream's files.
PR #32, flooring integer Divide functors	ml-explore#4515 and ml-explore#4642 add a FloorDivide primitive	The fork functors are removed, so integer Divide truncates, as in upstream. The PR #32 tests stay. The fork adds one guard to upstream's CPU floor_divide for min // -1.
PR #15, bias-sum operand promotion	ml-explore#4516 promotes the same operands	The sentence is removed from fork.yaml.

Conflicts and how I resolved them

git merge stopped on 69 files. 39 of them were files that only PR #10 changed, and PR #10 equals v0.32.2 there. For those 39 files I took upstream's file. They are listed at the end of this section.

PR #10 applied upstream content as a fork commit, so git used the old base 77a0c1e8 and saw that content as fork work. For the files that the fork changed after v0.32.2, I merged a second time with v0.32.2 as the base. That merge matches the backport lines in PR #27 to upstream's lines. I compared the two merges and checked each difference. A function-level audit then compared each definition in the result with the fork, upstream and v0.32.2.

File	Resolution
.github/actions/{build,build-wheel,test-linux,test-wheel,test-windows}/action.yml, .github/workflows/release.yml	The fork deleted them (PR #23). Upstream changed them. They stay deleted.
.github/workflows/update_bypass_list.yml, .github/scripts/{find-eligible-contributors,update-bypass-list}.js	No conflict. The merge added them again. I deleted them again (PR #23).
.github/workflows/build_and_test.yml	Upstream changed only the jobs that the fork deleted. Kept the fork file.
.gitignore, benchmarks/python/swiglu_bench.py, benchmarks/python/swiglu_qmm_bench.py	These came from the PR #27 backport. Took upstream.
mlx/backend/cpu/binary_ops.h and mlx/backend/metal/kernels/binary_ops.h	Took upstream FloorDivide. Removed the PR #32 Divide functors.
mlx/backend/cpu/simd/base_simd.h, mlx/backend/cpu/simd/accelerate_simd.h	Upstream's CPU floor_divide divides min by -1. That overflows and traps on x86: the PR #32 test crashed with SIGFPE in the Fedora x86_64 job. Added the PR #32 guard: divide by 1 there, so the result wraps to min.
tests/CMakeLists.txt	Upstream renamed METAL_TEST_SOURCES to GPU_TEST_SOURCES. Kept gptoss_mxfp4_tests.cpp under the new name.
python/tests/test_conv.py	Took upstream. Kept the GPU-only skip in test_conv2d_winograd_batch_tiling (PR #30). Added the same skip to upstream's new test_conv2d_winograd_disable, because the fork's CPU run cannot tell the two paths apart.
python/tests/test_fft.py, tests/ops_tests.cpp	Took upstream. The fork tests and skips merged clean.
python/tests/test_ops.py, python/tests/test_quantized.py, python/tests/test_fast_sdpa.py	Kept the fork tests next to upstream's new tests. test_fast_sdpa.py imports numpy again, because upstream removed the import and the fork tests use it.
mlx/backend/metal/device.h, mlx/fast.h	Context conflicts only. Took upstream's lines. The fork declarations merged clean.
mlx/memory.h	Kept both include lists.
mlx/compile.cpp, mlx/compile_impl.h	Kept the fork's compile-cache registry and the compatible compile_erase(fun_id) / compile_clear_cache(). The rest is upstream.
mlx/backend/metal/kernels/sdpa_vector.h	git merged it without a conflict, but the result had two copies of sdpa_vector_2pass_1_gqa. I rebuilt it from upstream and applied the FP32 partials again (PR #16).
mlx/backend/metal/kernels/scaled_dot_product_attention.metal, mlx/backend/metal/scaled_dot_product_attention.cpp	Took upstream's GQA, D512 and tk1 work. Renamed every two-pass kernel name, including the new gqa_<g> and _tk1 names, to sdpa_vector_2pass_fp32partials_*. Kept the float32 intermediate.
mlx/backend/metal/kernels/quantized.h	Kept the row and packed-column bound in the affine load_safe. The merge put upstream's new typename U = T template parameter on the fork's qmm_t_expert_impl. I moved it to qmm_t_impl.
mlx/backend/metal/kernels/quantized_nax.h	Kept the fork's distinct names and column bounds. Renamed get_pack_factor in upstream's new staged and split-K code. Kept the revert of the ml-explore#4352 simdgroup skip in affine_gather_qmm_rhs_nax.
mlx/backend/metal/kernels/fp_quantized.h	Took upstream's global scale and Fp4HalfLoader. Kept the fork names and column bounds. fp_gather_qmv_fast_tail now follows upstream's fp_gather_qmv_fast: the new partial_rows template parameter and the global scale in adjust_matrix_offsets.
mlx/backend/metal/kernels/fp_quantized_nax.h	Took upstream's global scale and the new fp_qmv_nax. Kept the fork names, the out-of-bounds scale guard and the revert of the ml-explore#4352 simdgroup skip in fp_gather_qmm_rhs_nax.
mlx/backend/metal/quantized.cpp	Took upstream's global scale, the fixed buffer slots and NAX for non-transposed calls (ml-explore#4684). Kept the fork's bool gather_qmm_rhs with the retract, the Gemma 4 expert-tile route, the takes_sorted_rhs_route gate, the GPT-OSS prefill tile and the decode fast tail.
mlx/backend/metal/primitives.cpp, mlx/fence.h, CLAUDE.md	git merged them without a conflict, but PR #10 content that upstream later removed came back: duplicate CrossEntropy Metal stubs (ml-explore#4520), the Fence::cast const overload (ml-explore#4552) and the CLAUDE.md symlink. Took upstream.
scripts/check_forkdiff.py	Upstream turned on ruff rule E501. Wrapped four long lines.
mlx/backend/cpu/gemm.h, gemms/cblas.cpp, masked_mm.cpp, conv.cpp	Latent upstream bug (worth an upstream issue; I did not open one). The fork CPU pass found NaN in upstream's test_fp_gather_qmm_non_multiple_of_32 (K=512, N=1040, sorted, float32). The NaN was only in the reference x @ w_sel (CPU matmul, M=1). Its 266240-byte output reused the buffer of the previous subtest's packed fp4 weights, and Accelerate read C with beta = 0 (NaN * 0). One helper, clear_gemm_output in gemm.h, now fills the exact C region with zeros before every BLAS gemm with beta = 0: the float32/float64/complex64 matmul wrappers (Matmul, AddMM, SegmentedMM), BlockMaskedMM and GatherMM in masked_mm.cpp, and the two explicit-gemm convolutions in conv.cpp. No other BLAS gemm or gemv call exists in the CPU backend. MLX's own SIMD gemm already skips C when beta is 0.
mlx/scheduler.{h,cpp}, mlx/transforms.cpp, tests/scheduler_tests.cpp	Upstream ordering race. StreamThread::enqueue throws the stream's stored task error on the main thread (ml-explore#3742). async_eval enqueues several tasks, so when the stream thread finishes a failing task first, async_eval throws. ml-explore#4338 expects the error only from synchronize. test_async_eval_error_in_synchronize failed on 2 of 6 heads in the fork CPU pass. async_eval now holds scheduler::DeferStreamErrors, so synchronize raises the error. A deterministic C++ test covers it.
mlx/ops.cpp (qqmm)	Upstream's CPU QQMatmul has no global-scale input and ignored the scales: wrong results with no error (upstream TODO at test_quantized.py:335). The fork CPU test pass found it. qqmm now raises "Global scale is only supported on the GPU" on a non-GPU stream, as gather_qmm does.
python/tests/test_fast_sdpa.py, python/tests/test_quantized.py (GPU-path tests)	Upstream tests that check only GPU kernels (test_sdpa_vector_2pass_threadgroup_limit with force_fused=True, the global-scale qqmm and gather_qmm tests) now check, when the default device is the CPU, that the documented error is raised. The GPU pass runs them as before.
python/tests/test_reduce.py (test_init_reduce_large_output)	Capability skip, ruled by David on 2026-10-23. The test skips only when the default device cannot hold its 4 GiB output: on the GPU when max_buffer_length is less than 4 GiB, on the CPU when memory_size is less than 16 GiB. The hosted macos-26 runner caps Metal buffers at 3.5 GiB. The test cannot be smaller: it is the regression test for ml-explore#4437 and needs 2^32 output elements (more than one 32-bit grid dimension), and bool outputs of one byte each give 4 GiB.
.github/workflows/fork-tests.yml	When a CPU test file fails, the job prints the end of its log at once.
tests/allocator_tests.cpp, tests/gptoss_mxfp4_tests.cpp	Upstream allocator::malloc now returns an owning allocator::Data, and allocator::free is gone. The fork tests use allocator::allocator().malloc / free, as upstream's tests do. Upstream gather_qmm has a new global_scale argument; the GPT-OSS tests pass std::nullopt.

Files that only PR #10 changed, now upstream's files (39): .pre-commit-config.yaml, CMakeLists.txt, CONTRIBUTING.md, setup.py, mlx/version.h, mlx/ops.cpp, mlx/primitives.cpp, mlx/fft.cpp, mlx/backend/metal/{conv.cpp,kernels.h,nojit_kernels.cpp}, mlx/backend/metal/kernels/fp_quantized_nax.metal, mlx/backend/metal/kernels/steel/attn/kernels/{steel_attention.metal,steel_attention_nax.h,steel_attention_nax.metal}, mlx/backend/cuda/{CMakeLists.txt,event.cu,fence.cpp}, mlx/backend/no_gpu/{fence.cpp,primitives.cpp}, python/src/{CMakeLists.txt,fast.cpp,indexing.cpp,ops.cpp,stream.cpp}, python/mlx/nn/losses.py, python/mlx/optimizers/optimizers.py, python/tests/{run.py,test_array.py,test_compile.py,test_eval.py,test_fast.py,test_load.py,test_reduce.py,test_vmap.py}, tests/load_tests.cpp, benchmarks/python/sdpa_bench.py, examples/cmake_extension/{pyproject.toml,requirements.txt}.

fork.yaml and FORKDIFF.md

• base.hash is now c6abf0dbb4fc141fad739b32ae486c06728e49e5, the merge-base with upstream main.
• Removed sections: "Upstream v0.32.2 sync carried in PR chore: upgrade core to MLX 0.32.2 #10" and "Upstream gather_mm and gather_qmm row-tile scheduling". Upstream has both changes.
• New sections "CPU qqmm refuses a global scale" (mlx/ops.cpp) and "CPU BLAS gemm clears C when beta is 0" (mlx/backend/cpu/{gemm.h,gemms/cblas.cpp,masked_mm.cpp,conv.cpp}). The coordinator added "64-bit products in the JACCL ring and the vector SDPA layout check" for two CodeQL findings.
• The simdgroup-skip revert has its own section now. Before, it was a paragraph in the removed backport section.
• "Integer floor_divide rounds down" is now "Integer floor_divide tests".
• "Fork tooling" lists the bypass-list workflow and scripts.
• FORKDIFF.md has a new sync policy. The fork merges upstream when David directs it. Each merge moves base.hash to the new merge-base and checks every section again. Between merges, an upstream fix that the fork needs is ported as its own commit.

python3 scripts/check_forkdiff.py --upstream-ref refs/remotes/upstream/main passes: 85 changed files, 85 described.

Downstream pins

Layr-Labs/mlx-c and Layr-Labs/mlx-swift point at the head SHA of this PR. Merge this PR with a merge commit, not a squash. A squash makes a new SHA, and the dependents must then pin again.

Test plan

I did not build or run tests on the laptop. The checks are:

• Fork Tests / metal-build-and-test: Metal build, C++ and Python tests, the second run on the CPU, the coverage report.
• Build and Test / Check Lint: pre-commit (clang-format, cmake-format, ruff).
• Build and Test / Linux Sanitizer Tests (ASAN) and (UBSAN).
• Build and Test / Linux Fedora (x86_64) and (aarch64).
• Fork Diff Check.
• A Release build on a Mac box (-DMLX_BUILD_TESTS=ON -DMLX_METAL_JIT=OFF) found compile errors in the merged kernels and in two fork test files. The fixes are in this PR. The box build of dbf463fc (library, mlx.metallib, C++ tests) has 0 errors. On the box I ran only the CPU test "test floor_divide integer edge cases on the CPU": it passes. The box GPU needs a lock protocol, so CI runs the GPU tests.

Observed CI flake

On head 1f7e9063 the ring distributed test in the GPU pass failed: the first failure is "[METAL] Command buffer execution failed: Caused GPU Timeout Error" in test_all_reduce_extra, and four more failures followed it. The same ring tests passed on 4588237c and aaec1b63. The only code change between aaec1b63 and 1f7e9063 is the qqmm guard, which the ring test does not use. I read it as a GPU timeout on the hosted runner. I did not change the test. Upstream's test-macos action does not fail the step on these ring failures; it checks only for [WARN].

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes (I ran the same hook versions on the changed files: clang-format 21.1.8, ruff 0.16.10, cmake-format 0.6.13.)
• I have added tests that prove my fix is effective or that my feature works (No new feature. The merge keeps the fork tests and adds upstream's tests.)
• I have updated the necessary documentation (if needed) (fork.yaml and FORKDIFF.md are updated. No other docs change.)
• I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: An AI agent (Claude Opus 5.5 in Claude Code) did the merge, resolved the conflicts, wrote the commits and wrote this description. David directed the work. The fork inherits upstream's rule against AI-written PR text, so the box above is not ticked. David decides.

🤖 Generated with Claude Code

Reference uniqueness without buffer ownership (commit 699649d38164)

Upstream 5dbe9991 (ml-explore#4634) makes is_donatable() require that MLX owns the buffer. That is right for donation. But the fork's mlx-c buffer-info API used is_donatable() as a "no other reference" proof, which mlx-swift-lm relies on for paged KV segments, checkpoint memory accounting and MTP state, including wrapped host pages. The fork adds array::is_uniquely_referenced(), the same reference-count test without the ownership test. is_donatable() and donation do not change. Pushed by the orchestrator; fork.yaml section "Reference uniqueness without buffer ownership".