ci: dynamically limit parallel build jobs to prevent OOM errors

[prasadn1]

Heavy C++ compilation (especially linking) can exhaust memory on GitHub Actions runners with high core counts but limited RAM, leading to intermittent OOM failures.

This introduces a cross-platform Python script that calculates a safe -j parallel build limit based on the system's available memory, ensuring builds scale safely across different runner types.

Proposed changes

Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.

No issue number. Pulled this from internal bug tracking

The Problem:
The current build configuration uses unbounded parallelism based on CPU core count (e.g., -j $(nproc) and %NUMBER_OF_PROCESSORS%). Because heavy C++ compilation and linking for ML kernels require significant
RAM per job (often 3-4GB per core), runners with high core counts but limited RAM (e.g., 8 cores / 16GB RAM) attempt to launch too many heavy processes simultaneously, exceeding the memory limit and causing
the build to crash.

The Solution:
This PR introduces a cross-platform Python utility (.github/scripts/set_cmake_parallel.py) that acts as a "defense in depth" system fix:

1. It dynamically queries the system's total physical memory.
2. It calculates a memory-safe parallel job limit (reserving ~3-4GB per core for Windows/MSVC and ~3GB for Linux/GCC).
3. It bounds the CMAKE_BUILD_PARALLEL_LEVEL environment variable.

This ensures the build system remains hardware-aware, protecting the current 16GB runners from OOMs while allowing the build to automatically and safely scale up if/when higher-spec 32GB or 64GB runners are
provisioned.

Checklist

Put an x in the boxes that apply.

• 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 have added tests that prove my fix is effective or that my feature works (Note: CI configuration changes are difficult to unit-test. I have verified the Darwin logic locally, but the Linux/Windows branches will need to be verified by the CI run on this PR).
• [] I have updated the necessary documentation (if needed)

[zcbenz]

We tried to limit parallel builds before but it made compilation time unbearable long. The special thing in our case is that we only have a few kernels that take insane RAM to build so using swap space ends up as a better solution for us:

mlx/.github/actions/setup-linux/action.yml

Lines 57 to 61 in 211e57b

	- name: Set swap space
	if: ${{ startsWith(inputs.toolkit, 'cuda') }}
	uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
	with:
	swap-size-gb: 16

[prasadn1]

If swap is the preferred stopgap, I can update this PR to expand the Windows pagefile during setup so it mirrors the Linux workaround. Alternatively, if you want a stricter build-system fix, we could use CMake Job Pools. We could define a heavy_compilation_pool limited to 1 or 2 jobs and explicitly assign only those specific memory-hungry kernels to it. That allows the other 95% of the framework to safely compile at -j$(nproc) while ensuring the heavy kernels serialize. This should prevent both OOMs and swap-thrashing without penalizing total CI time. Are either of these options helpful?

[zcbenz]

Didn't know CMake Job Pools before and it sounds a perfect solution! Most of the heavy kernels are under backend/cuda/quantized/qmm.

For Windows we are only doing CPU build though because the free runner does not have enough disk space for installing CUDA toolkits and buildings.

[prasadn1]

I've pivoted the PR to implement a strict build-system fix using CMake JOB_POOLS (supported by the Ninja generator)

1. I defined a heavy_compilation_pool with a concurrency limit of 2 in the root CMakeLists.txt
2. I explicitly assigned the generated .cu and .cpp targets in backend/cuda/quantized/qmm to this pool
   This should allow 95% of the framework to safely compile at -j$(nproc), keeping the build fast, while mathematically guaranteeing that the heaviest template instantiations serialize

Regarding windows: Do you happen to know which specific files/targets in the core C++ or CPU/Metal backends are causing the memory spikes on the Windows and ARM runners? If we can identify the heaviest targets there, we can assign them to this heavy_compilation_pool as well

[zcbenz]

Can you only enable job pool when the CI environment variable is true?

[prasadn1]

Done

[zcbenz]

Regarding windows: Do you happen to know which specific files/targets in the core C++ or CPU/Metal backends are causing the memory spikes on the Windows and ARM runners? If we can identify the heaviest targets there, we can assign them to this heavy_compilation_pool as well

The CPU/Metal backends are totally fine building on CI, this is only a CUDA kernel problem.