Repository navigation
Restore the upstream sync history squashed by #38 - #42
Merged
Merged
Conversation
…lore#4284) Co-authored-by: Cheng <[email protected]>
Co-authored-by: Daniel Hiltgen <[email protected]> Co-authored-by: Cheng <[email protected]>
…e#4356) Co-authored-by: Cheng <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]> Co-authored-by: Cheng <[email protected]>
Co-authored-by: Cheng <[email protected]>
Signed-off-by: Pradyot Ranjan <[email protected]> Co-authored-by: Pradyot Ranjan <[email protected]> Co-authored-by: Cheng <[email protected]>
Co-authored-by: Cheng <[email protected]>
Upstream's test_init_reduce_large_output (ml-explore#4437) makes a 4 GiB boolean output. The fork CI runs on a GitHub-hosted macOS runner with a 3.5 GiB Metal buffer limit and 7 GB of memory. On the GPU the allocation fails; on the CPU it competes with the other test files for memory. Skip the test when the default device cannot hold the output. Machines with more memory still run it. When a CPU test file fails, print the end of its log at once. The job of the previous run was stopped before the full logs printed, so the cause of the test_eval failure on the CPU is not known. Co-Authored-By: Claude Opus 5.5 <[email protected]>
The CPU QQMatmul kernel has no global scale input and ignores the scales that qqmm passes to it. qqmm with global scales on a CPU stream returned a wrong result with no error: in the fork CPU test pass, test_fp_qmv_fast_unaligned_n_global_scale differed by about 1000. qqmm now raises "Global scale is only supported on the GPU" on a stream that is not a GPU stream, as gather_qmm already does. Upstream tests that check only GPU kernels failed in the CPU pass: test_sdpa_vector_2pass_threadgroup_limit (force_fused=True), test_fp_qmv_fast_unaligned_n_global_scale, test_qqmm_global_scale_matrix_kernels and test_gather_qmm_global_scale. When the default device is the CPU, they now check that the documented error is raised. The GPU pass runs them as before. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Add one temporary step after the build in metal-build-and-test. It runs the failing case of test_fp_gather_qmm_non_multiple_of_32 (K=512, N=1040, transpose=False, sorted, float32) on the CPU with a fresh allocator, after a warm-up that fills the allocator cache with NaN buffers, and with the buffer cache off, and on the GPU. It prints NaN counts per expert for w, w_hat, y_q and y_hat. The step has continue-on-error, so the full test suite still runs. The next push removes it. State in fork.yaml when test_init_reduce_large_output skips, why it cannot be smaller (ml-explore#4437), and that David ruled on 2026-10-23 to keep the skip. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Remove the temporary diagnostic step from metal-build-and-test. Its reimplementation found no NaN in any mode. Add a temporary job, diag-cpu-nan, that builds the wheel and runs only the diagnostic, so its log is ready after the build. It runs the real test_fp_gather_qmm_non_multiple_of_32 alone, then after the tests that come before it, and an instrumented copy of the test body (only the failing case, and all cases in the test order, each evaluated as one graph). The findings also go to ::notice annotations. The next push removes the job. Co-Authored-By: Claude Opus 5.5 <[email protected]>
The tests import numpy through mlx_tests. The diagnostic job installed only the wheel, so every run stopped at the import. Co-Authored-By: Claude Opus 5.5 <[email protected]>
…ore#4737) Co-authored-by: Cheng <[email protected]>
Co-authored-by: Cheng <[email protected]>
test_fp_gather_qmm_non_multiple_of_32 (K=512, N=1040, sorted, float32) gave NaN on the CPU. The diagnostic job found the NaN only in the reference y_hat = x @ w_sel, a batched matmul with M = 1. Its inputs had no NaN. The test failed alone too, but only after the earlier subtests of the same test: y_hat is 266240 bytes, the same size as the packed fp4 weights of the previous subtest, so the output reused that buffer from the allocator cache. Accelerate read C even with beta = 0 and computed NaN * 0. Fill C with zeros before cblas_sgemm, cblas_dgemm and cblas_cgemm when beta is 0. MLX's own SIMD gemm already skips C when beta is 0. Remove the temporary diagnostic job. Co-Authored-By: Claude Opus 5.5 <[email protected]>
The previous commit cleared C only in the gemms/cblas.cpp wrappers. Four more call sites call cblas_sgemm directly with beta 0: BlockMaskedMM and GatherMM (the CPU MoE gathered matmul) in masked_mm.cpp, and the two explicit-gemm convolutions in conv.cpp. They can give the same NaN * 0 from a reused output buffer. Add one helper, clear_gemm_output in gemm.h, that fills the exact C region of a call (rows by columns, with its ldc) with zeros. Use it at all five call sites, the wrappers included. No other BLAS gemm or gemv call exists in the CPU backend. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Split the fork-tests job into parallel jobs. build-wheel and build-cpp build in parallel. build-cpp also runs the C++ tests on the GPU and on the CPU. python-gpu (3 shards) and python-cpu (4 shards) run the Python tests from the uploaded wheel. metal-build-and-test merges the coverage profiles of all jobs, writes the same coverage report, and fails when another job failed. python_shards.py checks the shard map: each python/tests/test*.py file must be in each pass exactly one time. Split files must partition their test IDs. The last job fails when a test ran in two shards. fork-build-cache restores and saves ccache on pull requests and on main, and keeps a cache of the compiled Metal kernels through an xcrun wrapper keyed on the preprocessed source, the arguments and the toolchain. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
…re#4738) Co-authored-by: Cheng <[email protected]>
Split test_fast_sdpa and test_quantized into named tests and rest entries in the CPU shards, so that no CPU shard has an entry longer than test_gated_delta_sequential. Put the entries in longest-first order. Save the ccache files and the Metal kernel cache only when the build had misses. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
The full --version text of clang and metal names the install folder and the OS, and it changed between three macos-26 runner images, so every job started from an empty cache. Key ccache on the clang major.minor version and fall back to a cache of another clang version; ccache keeps compiler_check=content. Key the Metal kernel cache on the metal compiler version and the kernels hash, and fall back only to a cache of the same metal compiler. Each Metal entry key holds the first line of the metal --version text and the SDK build. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
Upstream 5dbe999 (ml-explore#4634) makes is_donatable() also require that MLX owns the buffer, so the allocator never reuses a wrapped external buffer. The fork's mlx-c buffer-info API reported `unique` from is_donatable(), and mlx-swift-lm uses that as a "no other reference" proof for paged KV segments, checkpoint memory accounting and MTP state, including wrapped host pages, which now always read as not unique. Add is_uniquely_referenced(): the same reference-count test without the ownership test. is_donatable() and donation behavior do not change. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
test_eval.test_async_eval_error_in_synchronize (upstream ml-explore#4338) failed on some runs of the fork's CPU test pass: mx.async_eval raised the LU factorization error that the test expects only from mx.synchronize. Cause: StreamThread::enqueue throws the stream's stored error when it runs on the main thread (ml-explore#3742). async_eval enqueues several tasks. When the stream thread finishes the failing task before async_eval enqueues the next one, that enqueue throws. The order depends on thread timing, so the failure is intermittent. The fork does not change this path. async_eval now holds a scheduler::DeferStreamErrors guard. While it lives, enqueue on that thread does not throw the stored error, and synchronize throws it later. eval does not change. Add a C++ test that stores the error before async_eval, so it fails on every run without the fix. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
The hosted macOS runners are a shared pool. Pack the GPU pass into 2 shards and the CPU pass into 3 shards, balanced by the durations that run 39902948153 measured. A run now uses 8 macOS jobs, and at most 6 run at the same time (before: 10 jobs, 8 at the same time). Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
The preprocessed source of the NAX kernels names the Metal toolchain mount folder, which has a random suffix on each runner. The 10 NAX kernels missed the Metal cache in every warm build (41 of 51 hits). Remove the suffix from the text that the key hashes, and keep the toolchain version. CPU shard 0 was the slowest shard in the cold and warm runs of b72420d (651 s and 470 s of tests). Move test_quantized:rest to CPU shard 1 (221 s). Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
ci: build once and run the Metal and CPU test passes in parallel
Merge the 9 upstream commits after c6abf0d: ml-explore#4729, ml-explore#4730, ml-explore#4731, ml-explore#4728, ml-explore#4736, ml-explore#4737, ml-explore#4724, ml-explore#4738 and ml-explore#4744. The merge has no conflict. Every merged file gets exactly upstream's change. ml-explore#4728 widens the affine qmv_wide input-vector tile from 5 to 8. The fork has no qmv_wide change in this repository, and the GPT-OSS MXFP4 decode tail is an fp gather_qmv kernel that ml-explore#4728 does not touch. Each input vector still accumulates in its own registers, so the results do not change. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Set base.hash to the new merge-base after the merge of the 9 upstream commits. No fork section changes: the merge has no conflict, and each merged file gets exactly upstream's change. Co-Authored-By: Claude Opus 5.5 <[email protected]>
#38 was squash-merged. Main got the right tree but not the upstream ancestry, so the Fork Diff Check merge-base no longer equals base.hash, and the pinned sync head 83dbdb0 is reachable only from refs/pull/38/head. This merges the original sync head back. The tree does not change. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
jonathan308
added a commit
that referenced
this pull request
Oct 26, 2026
Brings in main at f78f0fa (2026-10-24): the upstream ml-explore/mlx sync to afc03c5 (#38, with its history restored by #42) and the fork fixes #40 and #43-#46. Conflicts and how they were resolved: - mlx/distributed/jaccl/lib/jaccl/mesh.cpp: main's per-wire queue setup (upstream ml-explore#4607, a second cable for the mesh) with this branch's COMPLETION_QUEUE_DEPTH for every completion queue. - mlx/distributed/jaccl/lib/jaccl/mesh_impl.h: main's stride parameters for all_gather and sum_scatter, with this branch's progress-guard calls and stage_send in every send path. - mlx/distributed/jaccl/lib/jaccl/ring_impl.h: this branch keeps its guarded member dispatch_wires (pool-wire errors are reported, the group closes after all wires return). Main moved an unguarded copy into threadpool.h and called it as dispatch_wires(pool_, ...); the ring call sites call the member again. The all_gather slice code is main's (identical to this branch's backport of upstream ml-explore#4443). - quantized_nax.h, fp_quantized_nax.h: main's kernels. Main already keeps the revert of the upstream ml-explore#4352 simdgroup skip (no sg_active) and adds the ragged-K tail bound of upstream ml-explore#4009. - fork.yaml: main's file, plus the JACCL description and section. The upstream backports this branch carried (ml-explore#4443, ml-explore#4557, ml-explore#4558, the scoped_env helper of ml-explore#4474) are now upstream ancestry, so their "not fork work" section and its references are gone. The JACCL section now says that a mesh with more than one wire has one guard per wire. Also: - python/tests/ring_test_distributed.py: main's file. The auto-merge had added test_strict_init_after_fallback twice. - tests/jaccl: the mesh all_gather and sum_scatter calls pass main's new stride argument (equal to the chunk size, as main's own all_reduce_scatter_gather does). Tests: tests/jaccl (simulated verbs, ASan and UBSan): 10 of 10 passed. check_forkdiff.py --upstream-ref upstream/main: 111 changed files, all described, base afc03c5 is the merge-base. clang-format 21: no changes. Co-Authored-By: Claude Opus 5.5 <[email protected]>
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.
#38 was squash-merged. Main got the correct tree, but not the upstream ancestry. Two things broke:
git merge-base main upstream/mainis the old base, notbase.hash.83dbdb04is reachable only fromrefs/pull/38/head. mlx-c main pins it inCMakeLists.txt.This PR merges the original sync head back into main. It changes no files: the tree is identical to main.
Merge with "Create a merge commit". A squash would drop the history again.
🤖 Generated with Claude Code