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 more CUDA architectures for PyPi package (#2427)

Browse Source 
* add cuda sm 90

* add more archs
This commit is contained in:
Awni Hannun
2025-07-28 12:35:15 -07:00
committed by GitHub
parent ab0e608862
commit 641be9463b
4 changed files with 13 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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) {
#if __CUDA_ARCH__ < 900
atomic_add_general(out, val);
#else
atomicAdd(out, val);
#endif
}
inline __device__ void atomic_add(__nv_bfloat16* out, __nv_bfloat16 val) {