[CUDA] Expose compilation targets via cuda.get_arch_list

docs/src/python/cuda.rst

@@ -7,3 +7,4 @@ CUDA
   :toctree: _autosummary
 
   is_available
+  get_arch_list

mlx/backend/cuda/CMakeLists.txt

@@ -181,10 +181,38 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
     set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
   endif()
 endif()
+# Resolve all/all-major/native aliases.
+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()
 message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${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}\"")
+    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 "$<$<COMPILE_LANGUAGE:CXX>: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)

mlx/backend/cuda/cuda.h

@@ -5,6 +5,7 @@
 #include <string>
 #include <unordered_map>
 #include <variant>
+#include <vector>
 
 #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<std::string> get_arch_list();
+
 /* Get information about a CUDA device. */
 MLX_API const
     std::unordered_map<std::string, std::variant<std::string, size_t>>&

mlx/backend/cuda/device_info.cpp

@@ -227,6 +227,10 @@ bool is_available() {
   return true;
 }
 
+std::vector<std::string> get_arch_list() {
+  return {MLX_CUDA_ARCH_LIST};
+}
+
 } // namespace cu
 
 } // namespace mlx::core

mlx/backend/cuda/no_cuda.cpp

@@ -11,6 +11,10 @@ bool is_available() {
   return false;
 }
 
+std::vector<std::string> get_arch_list() {
+  return {};
+}
+
 } // namespace cu
 
 namespace fast {

python/src/cuda.cpp

@@ -1,6 +1,8 @@
 // Copyright © 2023-2025 Apple Inc.
 
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/vector.h>
 
 #include "mlx/backend/cuda/cuda.h"
 
@@ -16,4 +18,17 @@ 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 CUDA
+          backend is not enabled.
+      )pbdoc");
 }

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 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, [])
+
     def test_device_info_default(self):
         info = mx.device_info()
         self.assertIsInstance(info, dict)