[davidtai]

• I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude Opus 5.5 (an AI agent) wrote this change and this description. The box above is not ticked. David decides.

Result

The Fork Tests workflow takes 13.3 min with warm caches and 21.1 min with cold caches. The old single job took 55 min. All tests run exactly one time in each pass (GPU pass and CPU pass), and the coverage report has the same rows and columns.

Base: sync/upstream-2026-10-22 at 36f043e (merged into this branch). Head: a8c2635. Measured run: 39921651047. Attempt 1 is cold, because all caches of this pull request were deleted before the push. Attempt 2 is warm.

Design

Job	macOS jobs	Work
build-wheel	1	Check the shard map. Build the Python wheel. Upload it.
build-cpp	1	Build the C++ library and tests. Run the C++ tests on the GPU, then on the CPU. Upload the C++ profile and the test binary.
python-gpu	2 shards	Install the wheel. Run one GPU shard, one entry at a time. Shard 0 also runs the distributed test and the two extension examples.
python-cpu	3 shards	Install the wheel. Run one CPU shard with DEVICE=cpu, three entries at a time.
metal-build-and-test	1	Check that no test ran in two shards. Merge the profiles of all jobs. Write the coverage report. Fail when another job did not succeed.

The two build jobs start at the same time. The five test shards start when the wheel is ready. A run uses 8 macOS jobs, and at most 6 run at the same time. Before this change, a run used 1 macOS job. The workflow keeps cancel-in-progress for pull requests, so a new push cancels the old run and frees its runners. Run 39902085502 shows this: its Python jobs were cancelled before they got a runner.

The steps of each job are the same commands as before: build-macos for the builds, test-macos for the distributed test and the extension examples, and the CPU harness for the CPU pass. The GPU Python pass ran python/tests/run.py in one process. It now runs one child process for each shard entry, with the same environment (MLX_ENABLE_TF32=0, MLX_ENABLE_CACHE_THRASHING_CHECK=0, the Metal debug variables). Each entry writes its own log. A failed entry prints the end of its log at once. The CPU pass keeps CPU_EXCLUDE (2 tests, unchanged).

Files:

• .github/workflows/fork-tests.yml: the jobs above.
• .github/actions/fork-build-cache/action.yml: restore and save of ccache and of the Metal kernel cache.
• .github/actions/setup/action.yml: restore only. New inputs ccache-fallback and install-macos-deps. New output ccache-save-key.
• .github/scripts/fork-tests/python-shards.json: the shard map.
• .github/scripts/fork-tests/python_shards.py: the guard and the shard runner.
• .github/scripts/fork-tests/metal-cache/xcrun: the Metal kernel cache.
• fork.yaml: description and globs for the new files.

Shard map

An entry is a whole test file (test_x), one test (test_x.Class.test_y), or the rest of a split file (test_x:rest). Four CPU files are split: test_fast_gated_delta, test_fast_sdpa and test_quantized, because one of their tests or the whole file takes longer than the other entries. The map is balanced by the per-entry durations that run 39902948153 measured.

Shard	Entries	Longest entry (cold / warm)	Test step (cold / warm)
GPU 0	30 files, plus the distributed test and the extension examples	test_quantized 224 s / 240 s	356 s / 385 s
GPU 1	6 files	test_fast_sdpa 90 s / 93 s	307 s / 290 s
CPU 0	15	test_gated_delta_simdgroup 525 s / 529 s	526 s / 531 s
CPU 1	16	test_gated_delta_sequential 346 s / 403 s	422 s / 455 s
CPU 2	15	test_gated_delta_nax 349 s / 372 s	350 s / 375 s

The CPU 0 shard has one test as its floor (test_gated_delta_simdgroup). A smaller shard is not possible without a change to that test. The speed of a hosted runner changes by up to 2.2 times between runs. For example, test_gated_delta_simdgroup took 278 s, 315 s, 470 s and 649 s in four runs.

The guard

python_shards.py stops the run in three places:

1. check (first step of build-wheel, and again in each shard): every python/tests/test*.py file must be in the map of each pass exactly one time. A file is either one whole-file entry, or named entries with at most one rest entry. An entry that names no file fails.
2. run (each shard): the named tests of a split file must match tests, no test can be in two entries, and a split file without a rest entry must name all its tests. The shard count in the workflow must equal the shard count in the map.
3. verify (in metal-build-and-test): each shard writes the IDs of the tests that started. The job fails when a shard result is missing or when a test ID is in two shards.

A broken copy of the map (a missing file, a duplicate file, a file that is whole and split, an unknown file) gave five ::error:: lines and exit code 1.

Tests started in each pass (from verify):

Pass	Before (run 39885643836)	Cold	Warm
C++ doctest, GPU	313	314	314
C++ doctest, CPU	313	314	314
Python, GPU	1004	1004	1004
Python, CPU	1002	1002	1002

The base merge added one C++ scheduler test, so the C++ count is 314. No test was added to or removed from the Python suite. No skip was added.

Caches

The organization blocks hendrikmuhs/ccache-action. The build jobs use Homebrew ccache and actions/cache at commit caa2961 (v5.1.0).

Cache	Key	Restore fallback	Size
ccache, wheel	ccache-fork-tests-wheel-FLAGS-clang21.0-macOS-ARM64-metal-SHA-RUN-ATTEMPT	same flags and clang version, then same flags and any clang version	76 MB
ccache, C++	ccache-fork-tests-cpp-FLAGS-clang21.0-...	as above	75 MB
Metal kernels	metal-32023.883-KERNELS-KIND-RUN-ATTEMPT	same metal version and kernels hash, then same metal version	41 MB (284 MB on disk)

• FLAGS is a hash of CFLAGS, CXXFLAGS, the deployment target and the build kind. KERNELS is hashFiles('mlx/backend/metal/kernels/**').
• ccache keeps compiler_check=content. An object from another compiler is a miss, never a wrong hit.
• The Metal cache script is first in PATH in the build jobs. It caches only xcrun -sdk macosx metal -c X.metal -o X.air and the .metallib link. All other calls go to /usr/bin/xcrun. The key of an .air file is a SHA-256 of the metal compiler version line, the SDK build, the arguments and the preprocessed source. The key of the .metallib is a SHA-256 of the same version data, the arguments and the hashes of the .air inputs. A cache from another metal version is never restored.
• The full --version text and the preprocessed line markers name the Metal toolchain mount folder, which has a random suffix on each runner (MetalToolchain-v17.6.109.0.oxYh3g). The keys use the version line only, and the script removes the suffix from the line markers. Before this fix, every runner image had a new key, and the 10 NAX kernels missed in every warm build.
• Pull requests and pushes to main save a cache only when the build had misses. A warm run saves nothing.

Hit rates (ccache -s and the Metal cache log):

Build	Cold ccache	Warm ccache	Cold Metal	Warm Metal
wheel	0 / 167 (0%)	167 / 167 (100%): 155 direct, 12 preprocessed	0 / 51	51 / 51
C++	0 / 162 (0%)	162 / 162 (100%): 150 direct, 12 preprocessed	0 / 51	51 / 51

Artifacts in one run: wheel 89 MB (uploaded with no compression, because a wheel is a zip file), C++ test binary 14 MB, 6 profile files 0.3 to 0.6 MB each, test ID lists under 6 KB each. Retention is 3 days.

Measured times

Before: run 39885643836 on 1f7e906, one job, 55 min. After: run 39921651047 on a8c2635.

Part	Before	After, cold	After, warm
Setup	1.8 min	in each job, 0.6 to 1.2 min	same
Build (wheel and C++)	18.3 min, one after the other	10.5 min, in parallel (wheel 9.6 min, C++ 8.2 min)	wheel 2.2 min, C++ 3.6 min with C++ tests; build steps 56 s and 47 s
GPU tests	13.2 min	2 shards: 9.8 and 5.6 min per job	10.4 and 5.5 min
CPU tests	22 min	C++ in build-cpp (23 s); 3 Python shards: 9.7, 7.8 and 6.6 min per job	9.7, 8.4 and 7.0 min
Coverage report	0.4 min	0.6 min	0.4 min
Total wall-clock	55 min	21.1 min	13.3 min
macOS jobs (most at the same time)	1 (1)	8 (6)	8 (6)

Other cold runs of this branch: 17.3 min (27f2a2c, 7 shards), 14.0 min (2268806, 7 shards). Cold time depends on the runner: the cold wheel build step took from 312 s to 576 s.

Coverage (C++ lines, Total row): before 73.36% (on 1f7e906), cold 73.40%, warm 73.39%. The report has the same rows, columns and text, plus one line that says the passes run in parallel jobs. The small changes come from the new base code and from tests with threads.

Open items

• With warm caches, the test shards set the wall-clock: GPU 0 (10.4 min, with the openmpi install, the distributed test and the extension examples) and CPU 0 (9.7 min, one 529 s test). A move of test_quantized from GPU 0 to GPU 1 can make the two GPU shards nearly equal (estimate: about 1 min less).
• Each pull request run with misses saves about 240 MB of cache. The repository cache limit is 10 GB, and GitHub removes the oldest caches first. The first push to main after this lands saves the main caches that all pull requests restore.
• The ring distributed test step keeps the upstream check if $(grep "\[WARN\]" stderr.log). This check does not fail on [WARN] lines (run 39885643836 printed test_donation: command not found). The step is unchanged here.

🤖 Generated with Claude Code