GPU discovery (#3055)

Co-authored-by: Angelos Katharopoulos <[email protected]>

docs/src/python/devices_and_streams.rst

@@ -17,3 +17,5 @@ Devices and Streams
    set_default_stream
    stream
    synchronize
+   device_count
+   device_info

mlx/backend/cpu/CMakeLists.txt

@@ -40,7 +40,7 @@ add_dependencies(mlx cpu_compiled_preamble)
 
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp

mlx/backend/cpu/available.cpp

@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cpu/available.h"
-
-namespace mlx::core::cpu {
-
-bool is_available() {
-  return true;
-}
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/available.h

@@ -1,9 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-namespace mlx::core::cpu {
-
-bool is_available();
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/device_info.cpp

@@ -0,0 +1,59 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cpu/device_info.h"
+
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#endif
+
+namespace mlx::core::cpu {
+
+namespace {
+
+// Get CPU architecture string
+std::string get_cpu_architecture() {
+#if defined(__aarch64__) || defined(__arm64__)
+  return "arm64";
+#elif defined(__x86_64__) || defined(_M_X64)
+  return "x86_64";
+#elif defined(__i386__) || defined(__i386) || defined(_M_IX86)
+  return "x86";
+#elif defined(__arm__) || defined(_M_ARM)
+  return "arm";
+#else
+  return "unknown";
+#endif
+}
+
+// Get CPU device name
+std::string get_cpu_name() {
+#ifdef __APPLE__
+  char model[256];
+  size_t len = sizeof(model);
+  if (sysctlbyname("machdep.cpu.brand_string", &model, &len, NULL, 0) == 0) {
+    return std::string(model);
+  }
+#endif
+  return get_cpu_architecture();
+}
+
+} // anonymous namespace
+
+bool is_available() {
+  return true;
+}
+
+int device_count() {
+  return 1;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static auto info =
+      std::unordered_map<std::string, std::variant<std::string, size_t>>{
+          {"device_name", get_cpu_name()},
+          {"architecture", get_cpu_architecture()}};
+  return info;
+}
+
+} // namespace mlx::core::cpu

mlx/backend/cpu/device_info.h

@@ -0,0 +1,28 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+namespace mlx::core::cpu {
+
+bool is_available();
+
+/**
+ * Get the number of available CPU devices.
+ *
+ * For CPU, always returns 1.
+ */
+int device_count();
+
+/**
+ * Get CPU device information.
+ *
+ * Returns a map with basic CPU device properties.
+ */
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index = 0);
+
+} // namespace mlx::core::cpu

mlx/backend/cuda/CMakeLists.txt

@@ -19,8 +19,8 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu

mlx/backend/cuda/cuda.cpp

@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/cuda.h"
-
-namespace mlx::core::cu {
-
-bool is_available() {
-  return true;
-}
-
-} // namespace mlx::core::cu

mlx/backend/cuda/cuda.h

@@ -2,6 +2,10 @@
 
 #pragma once
 
+#include <string>
+#include <unordered_map>
+#include <variant>
+
 #include "mlx/api.h"
 
 namespace mlx::core::cu {
@@ -9,4 +13,9 @@ namespace mlx::core::cu {
 /* Check if the CUDA backend is available. */
 MLX_API bool is_available();
 
+/* Get information about a CUDA device. */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
 } // namespace mlx::core::cu

mlx/backend/cuda/device_info.cpp

@@ -0,0 +1,232 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/cuda/cuda.h"
+
+#include <cuda_runtime.h>
+#include <dlfcn.h>
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+namespace mlx::core {
+
+namespace {
+
+// NVML dynamic loading for accurate memory reporting
+// (cudaMemGetInfo only sees current process)
+
+typedef int nvmlReturn_t;
+typedef struct nvmlDevice_st* nvmlDevice_t;
+struct nvmlMemory_t {
+  unsigned long long total;
+  unsigned long long free;
+  unsigned long long used;
+};
+
+struct NVMLState {
+  void* handle = nullptr;
+  nvmlReturn_t (*nvmlInit_v2)() = nullptr;
+  nvmlReturn_t (*nvmlDeviceGetHandleByUUID)(const char*, nvmlDevice_t*) =
+      nullptr;
+  nvmlReturn_t (*nvmlDeviceGetMemoryInfo)(nvmlDevice_t, nvmlMemory_t*) =
+      nullptr;
+};
+
+bool nvml_init(NVMLState& nvml) {
+#ifdef _WIN32
+  nvml.handle = dlopen("nvml.dll", RTLD_LAZY);
+  if (!nvml.handle) {
+    nvml.handle = dlopen(
+        "C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvml.dll", RTLD_LAZY);
+  }
+#else
+  nvml.handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY);
+#endif
+  if (!nvml.handle)
+    return false;
+
+  nvml.nvmlInit_v2 =
+      (decltype(nvml.nvmlInit_v2))dlsym(nvml.handle, "nvmlInit_v2");
+  nvml.nvmlDeviceGetHandleByUUID =
+      (decltype(nvml.nvmlDeviceGetHandleByUUID))dlsym(
+          nvml.handle, "nvmlDeviceGetHandleByUUID");
+  nvml.nvmlDeviceGetMemoryInfo = (decltype(nvml.nvmlDeviceGetMemoryInfo))dlsym(
+      nvml.handle, "nvmlDeviceGetMemoryInfo");
+
+  if (!nvml.nvmlInit_v2 || !nvml.nvmlDeviceGetHandleByUUID ||
+      !nvml.nvmlDeviceGetMemoryInfo) {
+    return false;
+  }
+  return nvml.nvmlInit_v2() == 0;
+}
+
+bool nvml_get_memory(
+    NVMLState& nvml,
+    const char* uuid,
+    size_t* free,
+    size_t* total) {
+  if (!nvml.handle)
+    return false;
+  nvmlDevice_t device;
+  if (nvml.nvmlDeviceGetHandleByUUID(uuid, &device) != 0)
+    return false;
+  nvmlMemory_t mem;
+  if (nvml.nvmlDeviceGetMemoryInfo(device, &mem) != 0)
+    return false;
+  *free = mem.free;
+  *total = mem.total;
+  return true;
+}
+
+std::string format_uuid(const cudaUUID_t& uuid) {
+  char buf[64];
+  snprintf(
+      buf,
+      sizeof(buf),
+      "GPU-%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+      (unsigned char)uuid.bytes[0],
+      (unsigned char)uuid.bytes[1],
+      (unsigned char)uuid.bytes[2],
+      (unsigned char)uuid.bytes[3],
+      (unsigned char)uuid.bytes[4],
+      (unsigned char)uuid.bytes[5],
+      (unsigned char)uuid.bytes[6],
+      (unsigned char)uuid.bytes[7],
+      (unsigned char)uuid.bytes[8],
+      (unsigned char)uuid.bytes[9],
+      (unsigned char)uuid.bytes[10],
+      (unsigned char)uuid.bytes[11],
+      (unsigned char)uuid.bytes[12],
+      (unsigned char)uuid.bytes[13],
+      (unsigned char)uuid.bytes[14],
+      (unsigned char)uuid.bytes[15]);
+  return buf;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info_impl(int device_index) {
+  // Static cache of device properties including UUID (needed for NVML lookup)
+  static auto all_devices = []() {
+    // Get device count
+    int count = 0;
+    cudaGetDeviceCount(&count);
+
+    // Collect info for all devices
+    struct DeviceInfo {
+      std::unordered_map<std::string, std::variant<std::string, size_t>> info;
+      std::string uuid;
+    };
+
+    std::vector<DeviceInfo> devices;
+
+    for (int i = 0; i < count; ++i) {
+      cudaDeviceProp prop;
+      cudaGetDeviceProperties(&prop, i);
+
+      DeviceInfo dev;
+      dev.info["device_name"] = std::string(prop.name);
+      dev.uuid = format_uuid(prop.uuid);
+      dev.info["uuid"] = dev.uuid;
+
+      // Architecture string (e.g., "sm_89")
+      char arch[16];
+      snprintf(arch, sizeof(arch), "sm_%d%d", prop.major, prop.minor);
+      dev.info["architecture"] = std::string(arch);
+
+      // PCI bus ID (domain:bus:device.function)
+      char pci_id[32];
+      snprintf(
+          pci_id,
+          sizeof(pci_id),
+          "%04x:%02x:%02x.0",
+          prop.pciDomainID,
+          prop.pciBusID,
+          prop.pciDeviceID);
+      dev.info["pci_bus_id"] = std::string(pci_id);
+
+      // Compute capability as size_t (to match Metal's variant type)
+      dev.info["compute_capability_major"] = static_cast<size_t>(prop.major);
+      dev.info["compute_capability_minor"] = static_cast<size_t>(prop.minor);
+
+      devices.push_back(std::move(dev));
+    }
+    return devices;
+  }();
+
+  // Initialize NVML once for fresh memory reads
+  static NVMLState nvml;
+  static bool nvml_initialized = nvml_init(nvml);
+
+  if (device_index < 0 ||
+      device_index >= static_cast<int>(all_devices.size())) {
+    static auto empty =
+        std::unordered_map<std::string, std::variant<std::string, size_t>>();
+    return empty;
+  }
+
+  // Return a copy with fresh memory info
+  // Using thread_local to avoid locks while keeping free_memory fresh
+  thread_local auto device_info_copy =
+      std::unordered_map<std::string, std::variant<std::string, size_t>>();
+
+  device_info_copy = all_devices[device_index].info;
+
+  // Get fresh memory info - try NVML first (system-wide), fallback to
+  // cudaMemGetInfo (process-level)
+  size_t free_mem, total_mem;
+
+  if (nvml_initialized &&
+      nvml_get_memory(
+          nvml,
+          all_devices[device_index].uuid.c_str(),
+          &free_mem,
+          &total_mem)) {
+    // NVML succeeded - use system-wide memory
+  } else {
+    // Fallback to cudaMemGetInfo (process-scoped)
+    int prev_device;
+    cudaGetDevice(&prev_device);
+    cudaSetDevice(device_index);
+    cudaMemGetInfo(&free_mem, &total_mem);
+    cudaSetDevice(prev_device);
+  }
+
+  device_info_copy["free_memory"] = free_mem;
+  device_info_copy["total_memory"] = total_mem;
+
+  return device_info_copy;
+}
+
+} // anonymous namespace
+
+namespace gpu {
+
+bool is_available() {
+  return true;
+}
+
+int device_count() {
+  int count = 0;
+  cudaGetDeviceCount(&count);
+  return count;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  return device_info_impl(device_index);
+}
+
+} // namespace gpu
+
+namespace cu {
+
+bool is_available() {
+  return true;
+}
+
+} // namespace cu
+
+} // namespace mlx::core

mlx/backend/cuda/eval.cpp

@@ -3,7 +3,6 @@
 #include "mlx/backend/gpu/eval.h"
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/available.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"
 
@@ -11,10 +10,6 @@
 
 namespace mlx::core::gpu {
 
-bool is_available() {
-  return true;
-}
-
 void new_stream(Stream s) {
   // Force initalization of CUDA, so CUDA runtime get destroyed at last.
   cudaFree(nullptr);

mlx/backend/gpu/available.h

@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/api.h"
-
-namespace mlx::core::gpu {
-
-MLX_API bool is_available();
-
-} // namespace mlx::core::gpu

mlx/backend/gpu/device_info.h

@@ -0,0 +1,36 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::gpu {
+
+MLX_API bool is_available();
+
+/**
+ * Get the number of available GPU devices.
+ */
+MLX_API int device_count();
+
+/**
+ * Get information about a GPU device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
+} // namespace mlx::core::gpu

mlx/backend/metal/CMakeLists.txt

@@ -102,6 +102,7 @@ target_sources(
   mlx
   PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp

mlx/backend/metal/allocator.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 #include "mlx/backend/metal/allocator.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/backend/metal/resident.h"
 #include "mlx/memory.h"
@@ -43,16 +44,17 @@ MetalAllocator::MetalAllocator()
           }),
       residency_set_(device_) {
   auto pool = metal::new_scoped_memory_pool();
-  auto memsize = std::get<size_t>(device_info().at("memory_size"));
+  const auto& info = gpu::device_info(0);
+  auto memsize = std::get<size_t>(info.at("memory_size"));
   auto max_rec_size =
-      std::get<size_t>(device_info().at("max_recommended_working_set_size"));
-  resource_limit_ = std::get<size_t>(device_info().at("resource_limit"));
+      std::get<size_t>(info.at("max_recommended_working_set_size"));
+  resource_limit_ = std::get<size_t>(info.at("resource_limit"));
   block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
   gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
   max_pool_size_ = block_limit_;
   device(mlx::core::Device::gpu)
       .set_residency_set(residency_set_.mtl_residency_set());
-  bool is_vm = std::get<std::string>(device_info().at("device_name")) ==
+  bool is_vm = std::get<std::string>(info.at("device_name")) ==
       "Apple Paravirtual device";
   if (is_vm) {
     return;
@@ -249,8 +251,8 @@ size_t get_memory_limit() {
   return metal::allocator().get_memory_limit();
 }
 size_t set_wired_limit(size_t limit) {
-  if (limit > std::get<size_t>(metal::device_info().at(
-                  "max_recommended_working_set_size"))) {
+  if (limit > std::get<size_t>(
+                  gpu::device_info(0).at("max_recommended_working_set_size"))) {
     throw std::invalid_argument(
         "[metal::set_wired_limit] Setting a wired limit larger than "
         "the maximum working set size is not allowed.");

mlx/backend/metal/device_info.cpp

@@ -0,0 +1,58 @@
+// Copyright © 2026 Apple Inc.
+
+#include <sys/sysctl.h>
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/metal.h"
+
+namespace mlx::core::gpu {
+
+bool is_available() {
+  return metal::is_available();
+}
+
+int device_count() {
+  return 1;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  auto init_device_info = []()
+      -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
+    auto pool = metal::new_scoped_memory_pool();
+    auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device();
+    auto name = std::string(raw_device->name()->utf8String());
+    auto arch = std::string(raw_device->architecture()->name()->utf8String());
+
+    size_t memsize = 0;
+    size_t length = sizeof(memsize);
+    sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
+
+    size_t rsrc_limit = 0;
+    sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
+    if (rsrc_limit == 0) {
+      rsrc_limit = 499000;
+    }
+
+    return {
+        {"device_name", name},
+        {"architecture", arch},
+        {"max_buffer_length", raw_device->maxBufferLength()},
+        {"max_recommended_working_set_size",
+         raw_device->recommendedMaxWorkingSetSize()},
+        {"memory_size", memsize},
+        {"resource_limit", rsrc_limit}};
+  };
+  static auto device_info_ = init_device_info();
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+
+  if (device_index == 0) {
+    return device_info_;
+  } else {
+    return empty;
+  }
+}
+
+} // namespace mlx::core::gpu

mlx/backend/metal/eval.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <memory>
 
-#include "mlx/backend/gpu/available.h"
 #include "mlx/backend/gpu/eval.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/utils.h"
@@ -10,10 +9,6 @@
 
 namespace mlx::core::gpu {
 
-bool is_available() {
-  return true;
-}
-
 void new_stream(Stream stream) {
   if (stream.device == mlx::core::Device::gpu) {
     metal::device(stream.device).new_queue(stream.index);

mlx/backend/metal/metal.cpp

@@ -1,8 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <memory>
 
-#include <sys/sysctl.h>
-
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/backend/metal/utils.h"
@@ -49,36 +47,4 @@ void stop_capture() {
   manager->stopCapture();
 }
 
-const std::unordered_map<std::string, std::variant<std::string, size_t>>&
-device_info() {
-  auto init_device_info = []()
-      -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
-    auto pool = new_scoped_memory_pool();
-    auto raw_device = device(default_device()).mtl_device();
-    auto name = std::string(raw_device->name()->utf8String());
-    auto arch = std::string(raw_device->architecture()->name()->utf8String());
-
-    size_t memsize = 0;
-    size_t length = sizeof(memsize);
-    sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
-
-    size_t rsrc_limit = 0;
-    sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
-    if (rsrc_limit == 0) {
-      rsrc_limit = 499000;
-    }
-
-    return {
-        {"device_name", name},
-        {"architecture", arch},
-        {"max_buffer_length", raw_device->maxBufferLength()},
-        {"max_recommended_working_set_size",
-         raw_device->recommendedMaxWorkingSetSize()},
-        {"memory_size", memsize},
-        {"resource_limit", rsrc_limit}};
-  };
-  static auto device_info_ = init_device_info();
-  return device_info_;
-}
-
 } // namespace mlx::core::metal

mlx/backend/no_cpu/CMakeLists.txt

@@ -1,6 +1,6 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/eval.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/encoder.cpp

mlx/backend/no_cpu/available.cpp

@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cpu/available.h"
-
-namespace mlx::core::cpu {
-
-bool is_available() {
-  return false;
-}
-
-} // namespace mlx::core::cpu

mlx/backend/no_cpu/device_info.cpp

@@ -0,0 +1,22 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cpu/device_info.h"
+
+namespace mlx::core::cpu {
+
+bool is_available() {
+  return false;
+}
+
+int device_count() {
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core::cpu

mlx/backend/no_gpu/CMakeLists.txt

@@ -1,6 +1,7 @@
 target_sources(
   mlx
   PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp

mlx/backend/no_gpu/device_info.cpp

@@ -0,0 +1,22 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/gpu/device_info.h"
+
+namespace mlx::core::gpu {
+
+bool is_available() {
+  return false;
+}
+
+int device_count() {
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core::gpu

mlx/backend/no_gpu/eval.cpp

@@ -2,15 +2,11 @@
 
 #include <stdexcept>
 
-#include "mlx/backend/gpu/available.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/backend/gpu/eval.h"
 
 namespace mlx::core::gpu {
 
-bool is_available() {
-  return false;
-}
-
 void new_stream(Stream) {}
 
 void eval(array&) {

mlx/device.cpp

@@ -1,9 +1,9 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2026 Apple Inc.
 
 #include <stdexcept>
 
-#include "mlx/backend/cpu/available.h"
-#include "mlx/backend/gpu/available.h"
+#include "mlx/backend/cpu/device_info.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/device.h"
 
 namespace mlx::core {
@@ -44,4 +44,29 @@ bool is_available(const Device& d) {
   return false;
 }
 
+int device_count(Device::DeviceType type) {
+  switch (type) {
+    case Device::cpu:
+      return cpu::device_count();
+    case Device::gpu:
+      return gpu::device_count();
+  }
+  // appease compiler
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(const Device& d) {
+  switch (d.type) {
+    case Device::cpu:
+      return cpu::device_info(d.index);
+    case Device::gpu:
+      return gpu::device_info(d.index);
+  }
+  // appease compiler
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
 } // namespace mlx::core

mlx/device.h

@@ -1,9 +1,13 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #pragma once
 
 #include "mlx/api.h"
 
+#include <string>
+#include <unordered_map>
+#include <variant>
+
 namespace mlx::core {
 
 struct MLX_API Device {
@@ -30,4 +34,23 @@ MLX_API bool operator!=(const Device& lhs, const Device& rhs);
 
 MLX_API bool is_available(const Device& d);
 
+/** Get the number of available devices for the given device type. */
+MLX_API int device_count(Device::DeviceType type);
+
+/**
+ * Get information about a device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(const Device& d = default_device());
+
 } // namespace mlx::core

mlx/mlx.h

@@ -4,7 +4,7 @@
 
 #include "mlx/array.h"
 #include "mlx/backend/cuda/cuda.h"
-#include "mlx/backend/gpu/available.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/compile.h"
 #include "mlx/device.h"

mlx/scheduler.cpp

@@ -1,7 +1,7 @@
 // Copyright © 2023 Apple Inc.
 
 #include "mlx/scheduler.h"
-#include "mlx/backend/gpu/available.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/backend/gpu/eval.h"
 
 namespace mlx::core {

python/src/device.cpp

@@ -1,9 +1,11 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #include <sstream>
 
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/device.h"
 #include "mlx/utils.h"
@@ -63,4 +65,34 @@ void init_device(nb::module_& m) {
       &mx::is_available,
       "device"_a,
       R"pbdoc(Check if a back-end is available for the given device.)pbdoc");
+  m.def(
+      "device_count",
+      &mx::device_count,
+      "device_type"_a,
+      R"pbdoc(
+      Get the number of available devices for the given device type.
+
+      Args:
+          device_type (DeviceType): The type of device to query (cpu or gpu).
+
+      Returns:
+          int: Number of devices.
+      )pbdoc");
+  m.def(
+      "device_info",
+      &mx::device_info,
+      nb::arg("d") = mx::default_device(),
+      R"pbdoc(
+      Get information about a device.
+
+      Returns a dictionary with device properties. Available keys depend
+      on the backend and device type. Common keys include ``device_name``,
+      ``architecture``, and ``total_memory`` (or ``memory_size``).
+
+      Args:
+          d (Device): The device to query (defaults to the default device).
+
+      Returns:
+          dict: Device information.
+      )pbdoc");
 }

python/src/metal.cpp

@@ -9,6 +9,7 @@
 #include <nanobind/stl/vector.h>
 
 #include "mlx/backend/metal/metal.h"
+#include "mlx/device.h"
 #include "mlx/memory.h"
 #include "python/src/small_vector.h"
 
@@ -90,21 +91,8 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Stop a Metal capture.
       )pbdoc");
-  metal.def(
-      "device_info",
-      &mx::metal::device_info,
-      R"pbdoc(
-      Get information about the GPU device and system settings.
-
-      Currently returns:
-
-      * ``architecture``
-      * ``max_buffer_size``
-      * ``max_recommended_working_set_size``
-      * ``memory_size``
-      * ``resource_limit``
-
-      Returns:
-          dict: A dictionary with string keys and string or integer values.
-      )pbdoc");
+  metal.def("device_info", []() {
+    DEPRECATE("mx.metal.device_info", "mx.device_info");
+    return mx::device_info(mx::Device(mx::Device::gpu, 0));
+  });
 }

python/tests/test_device.py

@@ -113,5 +113,38 @@ class TestStream(mlx_tests.MLXTestCase):
         self.assertEqual(a.item(), b.item())
 
 
+class TestDeviceInfo(mlx_tests.MLXTestCase):
+    def test_device_count(self):
+        cpu_count = mx.device_count(mx.cpu)
+        self.assertIsInstance(cpu_count, int)
+        self.assertEqual(cpu_count, 1)
+
+        gpu_count = mx.device_count(mx.gpu)
+        self.assertIsInstance(gpu_count, int)
+        self.assertGreaterEqual(gpu_count, 0)
+
+    def test_device_info_cpu(self):
+        info = mx.device_info(mx.cpu)
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+        self.assertTrue(len(info["device_name"]) > 0)
+        self.assertIn("architecture", info)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available")
+    def test_device_info_gpu(self):
+        gpu_count = mx.device_count(mx.gpu)
+        for i in range(gpu_count):
+            info = mx.device_info(mx.Device(mx.gpu, i))
+            self.assertIsInstance(info, dict)
+            self.assertIn("device_name", info)
+            self.assertTrue(len(info["device_name"]) > 0)
+            self.assertIn("architecture", info)
+
+    def test_device_info_default(self):
+        info = mx.device_info()
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+
+
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

python/tests/test_memory.py

@@ -54,7 +54,7 @@ class TestMemory(mlx_tests.MLXTestCase):
         old_limit = mx.set_wired_limit(0)
         self.assertEqual(old_limit, 1000)
 
-        max_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        max_size = mx.device_info(mx.gpu)["max_recommended_working_set_size"]
         with self.assertRaises(ValueError):
             mx.set_wired_limit(max_size + 10)