From 1aacb1a59e0b566b7e37849edca633ff588c3dac Mon Sep 17 00:00:00 2001 From: Clint Branham Date: Wed, 29 Jul 2026 16:32:31 -0500 Subject: [PATCH] [CUDA] Make the CUDA include roots configurable Two include roots are hardcoded or missing in the SwiftPM CUDA build, and neither can be pointed elsewhere. CCCL is hardcoded to the copy inside the CUDA toolkit, for both the `-I` and the `MLX_CCCL_DIR` define that NVRTC uses at runtime. MLX pins its own CCCL version, and when the toolkit ships a different one the mismatch surfaces as overload ambiguity in device code rather than 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. cudnn-frontend and CUTLASS are not provided at all. The CMake build fetches both via FetchContent, but 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 Adds `MLX_CCCL_DIR` to override the CCCL root (defaulting to the toolkit copy, so existing behaviour is unchanged) and `MLX_CUDA_INCLUDE_PATHS`, a colon-separated list of additional roots. Both are applied to the C++ sources via Package.swift and to `.cu` compilation via the CudaBuild plugin, so device and host code compile against the same headers. This keeps the roots external rather than vendoring cudnn-frontend and CUTLASS as submodules; that remains an option if their provisioning should be automatic rather than configured. Co-Authored-By: Claude Opus 5 --- Package.swift | 28 +++++++++++++++++++++++----- Plugins/CudaBuild/plugin.swift | 18 +++++++++++++++++- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/Package.swift b/Package.swift index e36e0d862..845ca6efe 100644 --- a/Package.swift +++ b/Package.swift @@ -4,6 +4,23 @@ import PackageDescription +/// Additional include roots for the CUDA build, colon-separated. +/// +/// The CMake build fetches cudnn-frontend and CUTLASS itself, but SwiftPM has no +/// equivalent, so `cudnn_utils.cpp` and the `quantized/qmm` kernels cannot find +/// their headers unless the roots are supplied here. +let extraCudaIncludePaths = + (Context.environment["MLX_CUDA_INCLUDE_PATHS"] ?? "") + .split(separator: ":") + .map(String.init) + .filter { !$0.isEmpty } + +/// CCCL headers used to compile device code and to JIT at runtime. +/// +/// Defaults to the copy shipped with the CUDA toolkit. MLX pins its own CCCL +/// version, so allow pointing at that copy instead when the two differ. +let cudaCcclDir = Context.environment["MLX_CCCL_DIR"] ?? "/usr/local/cuda/include/cccl" + let noMetalCmlxExcludes = [ // Exclude Metal backend files, but keep no_metal.cpp for stubs // "mlx/mlx/backend/metal/no_metal.cpp", @@ -115,11 +132,12 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu", ] + noMetalCmlxExcludes - cxxSettings = [ - .unsafeFlags(["-I/usr/local/cuda/include"]), - .unsafeFlags(["-I/usr/local/cuda/include/cccl"]), - .define("MLX_CCCL_DIR", to: "\"/usr/local/cuda/include/cccl\""), - ] + cxxSettings = + [ + .unsafeFlags(["-I/usr/local/cuda/include"]), + .unsafeFlags(["-I\(cudaCcclDir)"]), + .define("MLX_CCCL_DIR", to: "\"\(cudaCcclDir)\""), + ] + extraCudaIncludePaths.map { .unsafeFlags(["-I\($0)"]) } linkerSettings = [ .linkedLibrary("gfortran", .when(platforms: [.linux])), diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift index f7d12dd4c..9f2535f23 100644 --- a/Plugins/CudaBuild/plugin.swift +++ b/Plugins/CudaBuild/plugin.swift @@ -132,6 +132,22 @@ struct CudaBuild: BuildToolPlugin { } let stdArgs = settings.cppLanguageStandard.map { ["--std", $0] } ?? [] + // Include roots that live outside the package: the CCCL copy to compile + // device code against, plus anything named in MLX_CUDA_INCLUDE_PATHS + // (cudnn-frontend, CUTLASS). Kept in step with Package.swift, which + // applies the same roots to the C++ sources. + let environment = ProcessInfo.processInfo.environment + var externalIncludeRoots: [String] = [] + if let ccclDir = environment["MLX_CCCL_DIR"] { + externalIncludeRoots.append(ccclDir) + } + let extraRoots = environment["MLX_CUDA_INCLUDE_PATHS"] ?? "" + externalIncludeRoots += extraRoots.split(separator: ":").map(String.init) + let externalIncludeArgs: [String] = + externalIncludeRoots + .filter { !$0.isEmpty } + .flatMap { ["-I", $0] } + for inputFile in sourceCuFiles + generatedCuFiles { let outputCpp = URL(string: inputFile.relativePath, relativeTo: outputDir)! .deletingPathExtension().appendingPathExtension("cpp") @@ -143,7 +159,7 @@ struct CudaBuild: BuildToolPlugin { arguments: ["compile"] + verboseFlag + stdArgs + incrementalFlag + [ "--clangpp", clangUrl.url.path, "-I", sourceDir.path, - ] + headerSearchPathArgs + [ + ] + headerSearchPathArgs + externalIncludeArgs + [ inputFile.path, "-o", outputCpp.path, ],