From 30561229c7b3c3eb1d9ba17d316857a67c04b271 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Fri, 22 Aug 2025 14:39:43 -0700
Subject: [PATCH] Fix allocation bug in NCCL (#2530)

---
 mlx/backend/cuda/distributed.cu | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlx/backend/cuda/distributed.cu b/mlx/backend/cuda/distributed.cu
index 90eeacd7c..4d2658534 100644
--- a/mlx/backend/cuda/distributed.cu
+++ b/mlx/backend/cuda/distributed.cu
@@ -24,6 +24,7 @@ void AllReduce::eval_gpu(
       out.copy_shared_buffer(in);
       return {in, out};
     } else {
+      out.set_data(allocator::malloc(out.nbytes()));
       return {in, out};
     }
   };