[CUDA] Make the CUDA include roots configurable

[GoodOlClint]

Proposed changes

A follow-up to #413 ("Allowing SPM to compile on Linux with CUDA"): let the CUDA include roots be pointed somewhere other than the CUDA toolkit. Three headers the CUDA sources need are either hardcoded to the wrong copy or not provided at all, and none of them can currently be redirected.

CCCL is hardcoded to the toolkit copy

Package.swift pins both the -I and the MLX_CCCL_DIR define (which NVRTC uses to JIT at runtime) to /usr/local/cuda/include/cccl. MLX pins its own CCCL version via FetchContent, so when the toolkit ships a different one, device code fails on ambiguity rather than on anything naming a version:

mlx/backend/cuda/device/binary_ops.cuh(50): error: more than one instance of overloaded
function "cuda::std::fmod" matches the argument list

Observed with CUDA 13.0 (CCCL 3.0.1) against MLX's pinned CCCL 3.1.3. Because the define feeds the runtime JIT path too, a skewed CCCL is not only a compile-time problem.

cudnn-frontend and CUTLASS are not provided at all

The CMake build fetches both via FetchContent; SwiftPM has no equivalent, so the build fails on headers that are simply absent:

mlx/backend/cuda/cudnn_utils.h:9:10: fatal error: 'cudnn_frontend.h' file not found
mlx/backend/cuda/quantized/qmm/qmv.cu:9:10: fatal error: 'cute/numeric/numeric_types.hpp' file not found

cudnn_utils.cpp is compiled and cudnn is linked, so this is a gap rather than a deliberate exclusion.

The change

• MLX_CCCL_DIR overrides the CCCL root, defaulting to the toolkit copy — so an unconfigured build is byte-identical to today.
• MLX_CUDA_INCLUDE_PATHS is a colon-separated list of additional roots (cudnn-frontend, CUTLASS).

Both are applied on both sides: to the C++ sources via Package.swift, and to .cu compilation via the CudaBuild plugin, so host and device code compile against the same headers. The plugin adds only roots that were explicitly configured — nvcc already finds its own toolkit CCCL without an explicit -I, so an unconfigured build sees no new flags there either.

This keeps the roots external and configured rather than vendoring cudnn-frontend and CUTLASS as submodules. Vendoring would make provisioning automatic and is a reasonable alternative if you'd rather these not be the consumer's problem — but it is a much larger change, and the environment knobs are useful regardless (the CCCL skew above needs an override even once the other two are vendored). Happy to go that way instead.

Verification

On macOS (Xcode 26.6, Swift 6.3.3) — inert, as the settings apply only on the CUDA branch of the manifest:

• pre-commit run --all-files — clean.
• scripts/verify-docs.sh — passes.
• xcodebuild build-for-testing -scheme mlx-swift-Package -destination 'platform=macOS' — builds; CmlxTests and MLXTests pass.

On a real CUDA device (DGX Spark / GB10, sm_121, CUDA 13.0.88, Swift 6.3.3, aarch64, Ubuntu 24.04):

• swift build completed cleanly in 466.89s with MLX_CCCL_DIR pointed at MLX's pinned CCCL 3.1.3 and MLX_CUDA_INCLUDE_PATHS supplying the cudnn-frontend and CUTLASS roots — replacing an ad-hoc pile of -Xcxx flags entirely.
• All 95 .cu files recompiled, and the per-output configuration stamps recorded an include list containing CCCL 3.1.3 plus both extra roots — i.e. the configured roots are what nvcc actually used.
• Decoded Llama-3.2-1B-Instruct-4bit at 196.4 tok/s and gemma-4-12B-it-4bit at 23.0 tok/s, with output identical to the flag-pile workaround (194–195 / 22.9 tok/s).

Note CI does not cover the SwiftPM CUDA lane — linux_build_cmake_cuda builds via CMake, where FetchContent already supplies these headers — so none of this is exercised there.

Related

Two other follow-ups to #413 are open alongside this one, and all three touch Plugins/CudaBuild/plugin.swift. They are functionally independent, but whichever merges first will leave the others needing a trivial rebase — happy to reorder or stack them however you prefer:

• [CUDA] Build the CUDA argument lists in steps #442 — building the CUDA argument lists in steps (readability / type-check cost).
• [CUDA] Allow selecting nvcc's host compiler, and invalidate on change #443 — selecting nvcc's host compiler (+ invalidating on change).

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 — no test added: these are manifest and build-plugin settings that only take effect in a CUDA-enabled build, which has no CI lane here (the CUDA job builds via CMake). Evidenced on-device instead, per the stamped include lists above.
• I have updated the necessary documentation (if needed) — both knobs are documented where they are read in Package.swift; happy to add a README section on configuring a CUDA build if you'd like them surfaced there.

[davidkoski]

Change looks good to me, thanks!

[davidkoski]

Conflict needs resolving, also #443.