nccl dep + default for cuda (#2526)

2025-09-01 04:24:36 +08:00 · 2025-08-21 17:57:49 -07:00
parent 9392fc3f88
commit f93f87c802
3 changed files with 9 additions and 5 deletions
--- a/setup.py
+++ b/setup.py
@@ -297,6 +297,7 @@ if __name__ == "__main__":
                "nvidia-cublas-cu12==12.9.*",
                "nvidia-cuda-nvrtc-cu12==12.9.*",
                "nvidia-cudnn-cu12==9.*",
+                "nvidia-nccl-cu12",
            ]
        else:
            name = "mlx-cpu"