zcbenz在ml-explore/mlx项目中提交了一个关于CUDA后端的Pull Request #1983，目前处于工作进展中（WIP）。该请求旨在为项目添加CUDA支持，以提升计算性能。

这篇文章主要讨论了在MLX项目中添加CUDA后端的开发进展和相关优化。以下是主要内容总结：

1. CUDA后端的开发背景：

   • 开发者zcbenz正在为MLX项目添加CUDA后端，目前已经可以运行教程示例。
   • CUDA后端的主要优势包括支持统一内存（Unified Memory）以及NVIDIA硬件在学术和大规模计算中的广泛应用。

2. 构建和测试：

   • 开发者提供了构建和测试CUDA后端的命令，建议在Ubuntu 22.04和CUDA 11.6环境下进行测试。
   • 开发者还分享了一些开发过程中常用的构建命令。

3. 性能优化：

   • 开发者对logistic_regression示例进行了性能分析，发现了一些性能瓶颈，主要包括：
     • 内核启动和等待结果的时间较长。
     • cudaMalloc调用频繁，未来可以通过引入缓冲区缓存来减少这些调用。
     • cudaMemPrefetch调用延迟了内核执行，移除预取操作后性能显著提升。
   • 通过优化Event的实现、减少预取调用、批量清理临时数组等措施，性能从500 it/s提升到2100 it/s。

4. 未来优化方向：

   • 开发者计划进一步优化内核和CPU代码，特别是在所有操作符实现完成后进行。
   • 未来可能会引入JIT编译支持，以加快编译速度。

5. 与其他后端的兼容性：

   • 开发者讨论了如何将CUDA后端与ROCm（AMD的GPU计算平台）后端兼容，建议使用#define或代码转换工具（如HIPify）来实现。

6. Jetson设备的支持：

   • 有用户反馈在NVIDIA Jetson设备上构建CUDA后端时遇到问题，开发者表示未来会考虑为这些设备提供更好的支持。

7. 代码拆分与合并：

   • 由于代码不断重构，开发者计划将当前PR拆分为多个小PR，分别处理不同的操作符和后端基础设施。
   • 最终，该PR被拆分为多个小PR并合并到主分支中。

8. 性能对比：

   • 开发者对比了MLX与PyTorch在CUDA上的性能，发现MLX的内核执行时间与PyTorch相当，但启动延迟较高，未来会继续优化。

总结：这篇文章详细记录了CUDA后端在MLX项目中的开发进展、性能优化过程以及未来的改进方向。开发者通过多次优化显著提升了CUDA后端的性能，并计划在未来进一步改进和扩展其功能。

1. CUDA支持扩展

   • 评论1指出，这些补丁不仅支持Apple平台，还增加了对Linux平台（CUDA 12及SM 7.0以上）的支持。
   • 引用："this set of patches also adds support for 'Linux [platforms] with CUDA 12 and SM 7.0 (Volta) and up'."
   • 评论5提到，MLX实现了“一次编写，随处运行”，支持在Linux CUDA和Apple Silicon/Metal上运行。
   • 引用："Write MLX, run it on Linux CUDA, and Apple Silicon/Metal."

2. 对Nvidia的潜在威胁

   • 评论2认为，如果这种实现打开了其他厂商的CUDA支持大门，可能会对Nvidia构成威胁，并引发真正的竞争。
   • 引用："If this opens up the door to other implementations this could be a real threat to Nvidia."
   • 评论9提到，Strix Halo的性能已经让一些用户认为DGX Spark可能过于昂贵，而如果Apple设备能完美支持CUDA，可能会更具吸引力。
   • 引用："Flawless CUDA on Apple gear would make it really tempting."

3. 统一内存架构的疑问

   • 评论3和评论7讨论了MLX的统一内存架构（UMA）在CUDA上的实现问题，特别是如何在CPU和GPU之间处理内存共享和复制。
   • 引用："How does this work when one of the key features of MLX is using a unified memory architecture?"
   • 引用："For Apple silicon, the memory is anyway shared between GPU and CPU. But for CUDA, this is not the case."

4. 法律与技术细节的澄清

   • 评论4纠正了关于Google v Oracle案件的误解，指出重新实现API并不侵犯版权，并澄清当前实现并非CUDA的重新实现。
   • 引用："SCOTUS found that re-implementing an API does not infringe copyright."
   • 引用："Also apparently this is not a re-implementation of CUDA."

5. 开发者背景与期待

   • 评论8提到，这一实现来自Electron的创建者zcbenz，表达了对开发者背景的认可。
   • 引用："It's coming from zcbenz who created Electron among others."
   • 评论6则期待未来能支持Mac硬件上的Linux驱动。
   • 引用："Now do linux support / drivers for Mac hardware!"