zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

add cuda sm 90

Browse Source
This commit is contained in:
Awni Hannun
2025-07-27 08:05:23 -07:00
parent 9acec364c2
commit 4b30015872
2 changed files with 1 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/cuda/CMakeLists.txt
View File

@@ -108,7 +108,7 @@ endif()
# Compute capability 7 is required for synchronization between CPU/GPU with # Compute capability 7 is required for synchronization between CPU/GPU with
# managed memory. TODO: Add more architectures for potential performance gain. # managed memory. TODO: Add more architectures for potential performance gain.
set(MLX_CUDA_ARCHITECTURES set(MLX_CUDA_ARCHITECTURES
"70;80" "70;80;90"
CACHE STRING "CUDA architectures") CACHE STRING "CUDA architectures")
message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES

4
mlx/backend/cuda/device/atomic_ops.cuh
View File

@@ -49,11 +49,7 @@ inline __device__ void atomic_add(__half* out, __half val) {
} }
inline __device__ void atomic_add(complex64_t* out, complex64_t val) { inline __device__ void atomic_add(complex64_t* out, complex64_t val) {
#if __CUDA_ARCH__ < 900
atomic_add_general(out, val); atomic_add_general(out, val);
#else
atomicAdd(out, val);
#endif
} }
inline __device__ void atomic_add(__nv_bfloat16* out, __nv_bfloat16 val) { inline __device__ void atomic_add(__nv_bfloat16* out, __nv_bfloat16 val) {