[CUDA] Make the CUDA include roots configurable

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])),

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,
                     ],