From db9a3ad3dec3c8039dccaf087b2a30de03229f26 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Sat, 24 Oct 2026 11:36:41 -0700 Subject: [PATCH 1/2] [CUDA] Expose compilation targets via get_arch_list --- docs/src/python/cuda.rst | 1 + mlx/backend/cuda/CMakeLists.txt | 27 +++++++++++++++++++++++++++ mlx/backend/cuda/cuda.h | 4 ++++ mlx/backend/cuda/device_info.cpp | 18 ++++++++++++++++++ mlx/backend/cuda/no_cuda.cpp | 4 ++++ python/src/cuda.cpp | 16 ++++++++++++++++ python/tests/test_device.py | 10 ++++++++++ 7 files changed, 80 insertions(+) diff --git a/docs/src/python/cuda.rst b/docs/src/python/cuda.rst index 932d36b5ef..4ffefa1da3 100644 --- a/docs/src/python/cuda.rst +++ b/docs/src/python/cuda.rst @@ -7,3 +7,4 @@ CUDA :toctree: _autosummary is_available + get_arch_list diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 50be6f404f..89b0c54abb 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -185,6 +185,33 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") +set(_mlx_cuda_architectures "${MLX_CUDA_ARCHITECTURES}") +if(MLX_CUDA_ARCHITECTURES STREQUAL "all") + set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_ALL}") +elseif(MLX_CUDA_ARCHITECTURES STREQUAL "all-major") + set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR}") +elseif(MLX_CUDA_ARCHITECTURES STREQUAL "native") + set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_NATIVE}") +endif() + +set(_mlx_cuda_arch_list) +foreach(arch IN LISTS _mlx_cuda_architectures) + if(arch MATCHES "^([0-9]+[af]?)-(real|virtual)$") + if(CMAKE_MATCH_2 STREQUAL "real") + list(APPEND _mlx_cuda_arch_list "sm_${CMAKE_MATCH_1}") + else() + list(APPEND _mlx_cuda_arch_list "compute_${CMAKE_MATCH_1}") + endif() + elseif(arch MATCHES "^[0-9]+[af]?$") + list(APPEND _mlx_cuda_arch_list "sm_${arch}" "compute_${arch}") + else() + list(APPEND _mlx_cuda_arch_list "${arch}") + endif() +endforeach() +list(JOIN _mlx_cuda_arch_list "," _mlx_cuda_arch_list) +target_compile_definitions(mlx + PRIVATE MLX_CUDA_ARCH_LIST="${_mlx_cuda_arch_list}") + # Configure Windows CUDA DLL loading. if(WIN32) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/wddm.cpp) diff --git a/mlx/backend/cuda/cuda.h b/mlx/backend/cuda/cuda.h index 410bae4c63..d6b8b73fef 100644 --- a/mlx/backend/cuda/cuda.h +++ b/mlx/backend/cuda/cuda.h @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/api.h" @@ -13,6 +14,9 @@ namespace mlx::core::cu { /* Check if the CUDA backend is available. */ MLX_API bool is_available(); +/* Get the CUDA architectures MLX was compiled for. */ +MLX_API std::vector get_arch_list(); + /* Get information about a CUDA device. */ MLX_API const std::unordered_map>& diff --git a/mlx/backend/cuda/device_info.cpp b/mlx/backend/cuda/device_info.cpp index a4a15e2850..d5df35e3ac 100644 --- a/mlx/backend/cuda/device_info.cpp +++ b/mlx/backend/cuda/device_info.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -227,6 +228,23 @@ bool is_available() { return true; } +std::vector get_arch_list() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess) { + cudaGetLastError(); + return {}; + } + if (count == 0) { + return {}; + } + std::vector arch_list; + std::istringstream archs(MLX_CUDA_ARCH_LIST); + for (std::string arch; std::getline(archs, arch, ',');) { + arch_list.push_back(std::move(arch)); + } + return arch_list; +} + } // namespace cu } // namespace mlx::core diff --git a/mlx/backend/cuda/no_cuda.cpp b/mlx/backend/cuda/no_cuda.cpp index 175a505b4b..8126790dd4 100644 --- a/mlx/backend/cuda/no_cuda.cpp +++ b/mlx/backend/cuda/no_cuda.cpp @@ -11,6 +11,10 @@ bool is_available() { return false; } +std::vector get_arch_list() { + return {}; +} + } // namespace cu namespace fast { diff --git a/python/src/cuda.cpp b/python/src/cuda.cpp index 13b3a0154b..38aff80db9 100644 --- a/python/src/cuda.cpp +++ b/python/src/cuda.cpp @@ -1,6 +1,8 @@ // Copyright © 2023-2025 Apple Inc. #include +#include +#include #include "mlx/backend/cuda/cuda.h" @@ -16,4 +18,18 @@ void init_cuda(nb::module_& m) { R"pbdoc( Check if the CUDA back-end is available. )pbdoc"); + + cuda.def( + "get_arch_list", + &mx::cu::get_arch_list, + R"pbdoc( + Get the CUDA architectures MLX was compiled for. + + Architectures use NVIDIA's ``sm_*`` (binary) and ``compute_*`` (PTX) + notation, for example ``["sm_80", "sm_120a", "compute_120"]``. + + Returns: + list(str): The compiled architectures, or an empty list if no + CUDA device can be used. + )pbdoc"); } diff --git a/python/tests/test_device.py b/python/tests/test_device.py index bef788bb1d..539c59b27b 100644 --- a/python/tests/test_device.py +++ b/python/tests/test_device.py @@ -184,6 +184,16 @@ def test_device_info_gpu(self): self.assertTrue(len(info["device_name"]) > 0) self.assertIn("architecture", info) + def test_cuda_get_arch_list(self): + arch_list = mx.cuda.get_arch_list() + self.assertIsInstance(arch_list, list) + if not mx.cuda.is_available() or mx.device_count(mx.gpu) == 0: + self.assertEqual(arch_list, []) + return + self.assertTrue(arch_list) + for arch in arch_list: + self.assertRegex(arch, r"^(sm|compute)_[0-9]+[af]?$") + def test_device_info_default(self): info = mx.device_info() self.assertIsInstance(info, dict) From 543ce631d6d3e0edf256678280beb9eedd21fa35 Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 27 Oct 2026 00:59:23 -0700 Subject: [PATCH 2/2] nit --- mlx/backend/cuda/CMakeLists.txt | 35 ++++++++++++++++---------------- mlx/backend/cuda/device_info.cpp | 16 +-------------- python/src/cuda.cpp | 5 ++--- python/tests/test_device.py | 10 ++++----- 4 files changed, 26 insertions(+), 40 deletions(-) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 89b0c54abb..0013892714 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -181,36 +181,37 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES) set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a") endif() endif() -message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") -set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES - "${MLX_CUDA_ARCHITECTURES}") - -set(_mlx_cuda_architectures "${MLX_CUDA_ARCHITECTURES}") +# Resolve all/all-major/native aliases. if(MLX_CUDA_ARCHITECTURES STREQUAL "all") - set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_ALL}") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_ALL}") elseif(MLX_CUDA_ARCHITECTURES STREQUAL "all-major") - set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR}") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR}") elseif(MLX_CUDA_ARCHITECTURES STREQUAL "native") - set(_mlx_cuda_architectures "${CMAKE_CUDA_ARCHITECTURES_NATIVE}") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_NATIVE}") endif() +message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") +set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES + "${MLX_CUDA_ARCHITECTURES}") -set(_mlx_cuda_arch_list) -foreach(arch IN LISTS _mlx_cuda_architectures) +# Pass the list of cuda arch to code. +set(MLX_CUDA_ARCH_LIST) +foreach(arch IN LISTS MLX_CUDA_ARCHITECTURES) if(arch MATCHES "^([0-9]+[af]?)-(real|virtual)$") if(CMAKE_MATCH_2 STREQUAL "real") - list(APPEND _mlx_cuda_arch_list "sm_${CMAKE_MATCH_1}") + list(APPEND MLX_CUDA_ARCH_LIST "\"sm_${CMAKE_MATCH_1}\"") else() - list(APPEND _mlx_cuda_arch_list "compute_${CMAKE_MATCH_1}") + list(APPEND MLX_CUDA_ARCH_LIST "\"compute_${CMAKE_MATCH_1}\"") endif() elseif(arch MATCHES "^[0-9]+[af]?$") - list(APPEND _mlx_cuda_arch_list "sm_${arch}" "compute_${arch}") + list(APPEND MLX_CUDA_ARCH_LIST "\"sm_${arch}\"" "\"compute_${arch}\"") else() - list(APPEND _mlx_cuda_arch_list "${arch}") + list(APPEND MLX_CUDA_ARCH_LIST "\"${arch}\"") endif() endforeach() -list(JOIN _mlx_cuda_arch_list "," _mlx_cuda_arch_list) -target_compile_definitions(mlx - PRIVATE MLX_CUDA_ARCH_LIST="${_mlx_cuda_arch_list}") +list(JOIN MLX_CUDA_ARCH_LIST "," MLX_CUDA_ARCH_LIST) +target_compile_definitions( + mlx + PRIVATE "$<$:MLX_CUDA_ARCH_LIST=${MLX_CUDA_ARCH_LIST}>") # Configure Windows CUDA DLL loading. if(WIN32) diff --git a/mlx/backend/cuda/device_info.cpp b/mlx/backend/cuda/device_info.cpp index d5df35e3ac..43c25ee8cb 100644 --- a/mlx/backend/cuda/device_info.cpp +++ b/mlx/backend/cuda/device_info.cpp @@ -6,7 +6,6 @@ #include #include -#include #include #include #include @@ -229,20 +228,7 @@ bool is_available() { } std::vector get_arch_list() { - int count = 0; - if (cudaGetDeviceCount(&count) != cudaSuccess) { - cudaGetLastError(); - return {}; - } - if (count == 0) { - return {}; - } - std::vector arch_list; - std::istringstream archs(MLX_CUDA_ARCH_LIST); - for (std::string arch; std::getline(archs, arch, ',');) { - arch_list.push_back(std::move(arch)); - } - return arch_list; + return {MLX_CUDA_ARCH_LIST}; } } // namespace cu diff --git a/python/src/cuda.cpp b/python/src/cuda.cpp index 38aff80db9..34dec0fc69 100644 --- a/python/src/cuda.cpp +++ b/python/src/cuda.cpp @@ -18,7 +18,6 @@ void init_cuda(nb::module_& m) { R"pbdoc( Check if the CUDA back-end is available. )pbdoc"); - cuda.def( "get_arch_list", &mx::cu::get_arch_list, @@ -29,7 +28,7 @@ void init_cuda(nb::module_& m) { notation, for example ``["sm_80", "sm_120a", "compute_120"]``. Returns: - list(str): The compiled architectures, or an empty list if no - CUDA device can be used. + list(str): The compiled architectures, or an empty list if CUDA + backend is not enabled. )pbdoc"); } diff --git a/python/tests/test_device.py b/python/tests/test_device.py index 539c59b27b..86f7fbebd1 100644 --- a/python/tests/test_device.py +++ b/python/tests/test_device.py @@ -187,12 +187,12 @@ def test_device_info_gpu(self): def test_cuda_get_arch_list(self): arch_list = mx.cuda.get_arch_list() self.assertIsInstance(arch_list, list) - if not mx.cuda.is_available() or mx.device_count(mx.gpu) == 0: + if mx.cuda.is_available(): + self.assertTrue(arch_list) + for arch in arch_list: + self.assertRegex(arch, r"^(sm|compute)_[0-9]+[af]?$") + else: self.assertEqual(arch_list, []) - return - self.assertTrue(arch_list) - for arch in arch_list: - self.assertRegex(arch, r"^(sm|compute)_[0-9]+[af]?$") def test_device_info_default(self): info = mx.device_info()