mlx-swift with CUDA for Linux and Windows?

[zamderax]

Hi MLX Team,

First off, fantastic work on mlx and mlx-swift! The recent addition of the CUDA backend to the core mlx library is a huge boon! I'm opening this issue to propose and discuss the possibility of extending mlx-swift to support this new CUDA backend on non-Apple platforms, specifically Linux and Windows. The Swift language has made significant strides in cross-platform support for Linux and Windows, there seems to be a clear opportunity to position mlx-swift as a truly powerful, cross-platform ML framework. This would allow developers to write ML code once in Swift and deploy it not only on Apple hardware but also on NVIDIA-powered servers, workstations, and embedded devices.

Does anyone know what it would take to create a build configuration or pathway for mlx-swift that links against the CUDA-enabled version of mlx? The primary targets would be:

• Linux (x86_64): For cloud servers and workstations.
• Linux (aarch64): Specifically for the NVIDIA Jetson family (Jetson Orin Nano, AGX, etc.).
• Windows (x86_64): For developers on Windows machines with NVIDIA GPUs.

[davidkoski]

The current version of mlx that mlx-swift builds against doesn't have the CUDA backend, but when we integrate it, we will do something like this:

• https://github.com/ml-explore/mlx-swift/blob/main/Package.swift#L76-L78

It isn't clear (yet) how we would allow this to build with a CUDA backend. There is a new SwiftPM feature that would probably enable it (variants):

• https://forums.swift.org/t/pitch-replaceable-library-plugins/77605

but it would require moving to a newer SwiftPM (and thus Xcode and maybe swift version).

There might be differences in Foundation as well that would need to be dealt with, though that seems to be getting better (converging) over time.

[zamderax]

What about using:

Swift PackageTrait to https://github.com/swiftlang/swift-evolution/blob/main/proposals/0450-swiftpm-package-traits.md

This was introduced in Swift 6.1

// Inside the targets
.target(
    name: "Cmlx",
    // Base exclusions that apply to all backends
    exclude: [
        "mlx-c/examples",
        "mlx/tests",
        "mlx/python",
        "mlx/examples",
        // ... other common excludes
    ],
    
    cxxSettings: [
        // Common header search paths
        .headerSearchPath("mlx"),
        .headerSearchPath("mlx-c"),

        // --- Default (Metal) Settings ---
        .define("_METAL_"),
        .define("MLX_USE_ACCELERATE"),

        // --- CUDA-specific Settings ---
        .define("_CUDA_", .when(trait: "cuda")),
        .define("MLX_USE_CUDA_", .when(trait: "cuda")) // Assuming a macro like this
    ],
    
    linkerSettings: [
        // --- Default (Metal) Linker Settings ---
        .linkedFramework("Foundation"),
        .linkedFramework("Metal"),
        .linkedFramework("Accelerate"),

        // --- CUDA-specific Linker Settings ---
        // These would be ignored on non-Linux platforms
        .linkedLibrary("cudart", .when(trait: "cuda")),
        .linkedLibrary("cublas", .when(trait: "cuda")),
        .linkedLibrary("nvrtc", .when(trait: "cuda"))
    ]
),

[davidkoski]

Yes, this is exactly what I was referring to (though it looked like it was Variant in the doc I found, but that was the one I meant!), but we still support older swiftpm versions so I don't think we can adopt this until we move forward off of them.

[zamderax]

1. Okay so I'll just create a fork, and make a PR and just keep it up to date from main?
2. What's your thought process of supporting older Swift PMs as 6.2 is right around the corner and it seems quite critical in supporting Swift Concurrency and Cross Platform Artifact Bundles for Linux, Windows and Embedded (static libraries not dynamic)

[zamderax]

Also I just had an idea

Option 1 (requires SPM 6.1 and higher)

1. use a Swift Package Trait called cuda like so:

traits: [
   .trait(name: "cuda")
],

and in linkerSettings

            linkerSettings: [
                .linkedFramework("Foundation"),
                .linkedFramework("Metal"),
                .linkedFramework("Accelerate"),
                // Conditionally link CUDA libraries when cuda trait is enabled
                .linkedLibrary("cudart", .when(traits: ["cuda"])),
                .linkedLibrary("cublas", .when(traits: ["cuda"])),
                .linkedLibrary("cufft", .when(traits: ["cuda"])),
                .linkedLibrary("cudnn", .when(traits: ["cuda"])),
            ]

and in the target exlude:

targets: [
        .target(
            name: "Cmlx",
            exclude: [
                // vendor docs
                "metal-cpp.patch",
                "vendor-README.md",

                // example code + mlx-c distributed
                "mlx-c/examples",
                "mlx-c/mlx/c/distributed.cpp",
                "mlx-c/mlx/c/distributed_group.cpp",

                // vendored library, include header only
                "json",

                // vendored library
                "fmt/test",
                "fmt/doc",
                "fmt/support",
                "fmt/src/os.cc",
                "fmt/src/fmt.cc",

                // Backend files are included via backend_selector.cpp
                "mlx/mlx/backend/no_cpu/compiled.cpp",
                "mlx/mlx/backend/cpu/compiled.cpp",
                "mlx/mlx/backend/cuda/compiled.cpp",
                "mlx/mlx/backend/cuda/no_cuda.cpp",
                "mlx/mlx/backend/no_gpu",


                // mlx files that are not part of the build
                "mlx/ACKNOWLEDGMENTS.md",
                "mlx/CMakeLists.txt",
                "mlx/CODE_OF_CONDUCT.md",
                "mlx/CONTRIBUTING.md",
                "mlx/LICENSE",
                "mlx/MANIFEST.in",
                "mlx/README.md",
                "mlx/benchmarks",
                "mlx/cmake",
                "mlx/docs",
                "mlx/examples",
                "mlx/mlx.pc.in",
                "mlx/pyproject.toml",
                "mlx/python",
                "mlx/setup.py",
                "mlx/tests",

                // opt-out of these backends (using metal)
                "mlx/mlx/backend/no_metal",
                "mlx/mlx/backend/no_cpu",

                // build variants (we are opting _out_ of these)
                "mlx/mlx/io/no_safetensors.cpp",
                "mlx/mlx/io/gguf.cpp",
                "mlx/mlx/io/gguf_quants.cpp",

                // see PrepareMetalShaders -- don't build the kernels in place
                "mlx/mlx/backend/metal/kernels",
                "mlx/mlx/backend/metal/nojit_kernels.cpp",

                // do not build distributed support (yet)
                "mlx/mlx/distributed/mpi/mpi.cpp",
                "mlx/mlx/distributed/ring/ring.cpp",

                // bnns instead of simd (accelerate)
                "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp",
                "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp",
            ],
            sources: [
                ".",
            ],

            cSettings: [
                .headerSearchPath("mlx"),
                .headerSearchPath("mlx-c"),
            ],

            cxxSettings: [
                .headerSearchPath("mlx"),
                .headerSearchPath("mlx-c"),
                .headerSearchPath("metal-cpp"),
                .headerSearchPath("json/single_include/nlohmann"),
                .headerSearchPath("fmt/include"),

                .define("MLX_USE_ACCELERATE"),
                .define("ACCELERATE_NEW_LAPACK"),
                .define("_METAL_"),
                .define("SWIFTPM_BUNDLE", to: "\"mlx-swift_Cmlx\""),
                .define("METAL_PATH", to: "\"default.metallib\""),
                .define("MLX_VERSION", to: "\"0.27.1\""),
                // Define MLX_BUILD_CUDA when cuda trait is enabled
                .define("MLX_BUILD_CUDA", .when(traits: ["cuda"])),
                // Define MLX_BUILD_CPU when cuda trait is NOT enabled
                .define("MLX_BUILD_CPU", .when(traits: [])),
            ],
            linkerSettings: [
                .linkedFramework("Foundation"),
                .linkedFramework("Metal"),
                .linkedFramework("Accelerate"),
                // Conditionally link CUDA libraries when cuda trait is enabled
                .linkedLibrary("cudart", .when(traits: ["cuda"])),
                .linkedLibrary("cublas", .when(traits: ["cuda"])),
                .linkedLibrary("cufft", .when(traits: ["cuda"])),
                .linkedLibrary("cudnn", .when(traits: ["cuda"])),
            ]
        ),

Option 2 - Create two targets

One Target Per Platform. "metal", "cuda", maybe in the future rocm etc... This but this makes the downstream user have more complications.

[davidkoski]

1. Okay so I'll just create a fork, and make a PR and just keep it up to date from main?

I think that works to start -- perhaps it will guide the feature once we are ready to take it back into the main repo.

2. What's your thought process of supporting older Swift PMs as 6.2 is right around the corner and it seems quite critical in supporting Swift Concurrency and Cross Platform Artifact Bundles for Linux, Windows and Embedded (static libraries not dynamic)

I think we should support current versions of Swift and SDKs (where possible) -- there are availability checks, etc. that we can use. I am not sure how to do the same with SwiftPM, but I agree it would be good to take advantage of newer versions if we could.

I do remember seeing this:

• https://github.com/apple/swift-numerics/tree/e0ec0f5f3af6f3e4d5e7a19d2af26b481acb6ba8

Notice the two versions of the Package.swift? I need to see what that feature is and how it works -- perhaps this is the way forward.

[davidkoski]

And that said, we can deprecate old versions but we should do so in a deliberate way.

[zamderax]

Got it

Regarding multiple Package.swift versionsL

https://github.com/swiftlang/swift-evolution/blob/main/proposals/0135-package-manager-support-for-differentiating-packages-by-swift-version.md

Looks like we have a path forward?

[davidkoski]

Yeah, I think we should be able to use that!

[zamderax]

I'll make a PR for you tonight

[zamderax]

I created a draft pr here: #259

Let me know if you like the direction so far.

[davidkoski]

#273 will get us to a version of mlx that has CUDA. I think one approach we could take is to conditionalize the existing Package.swift based on the OS. Build for CUDA on linux (and windows).

That said, I see these hurdles:

• we need to build the JIT versions of the kernels
  • this could be done the same way we do the metal versions
• we need to build some of the kernels with the cuda compiler
  • plugins? this produces a resource that we load at runtime so we need the same techniques to load the file that we have for the metallib
  • or maybe we have a JIT-only version? @awni another use case for this
• the build requires specific versions of CCCL, NVTX (and maybe others), plus various libraries installed on the system
  • swiftpm doesn't have a way to fetch these at build time so we would have to somehow document (and check?) for them at build time?
• we might need various env variables set, e.g. CUDNN_FRONTEND_SKIP_JSON_LIB

I think we could do all of this without traits as we don't need Metal and CUDA building for the same OS.