From 31c6f6e33f9dd612ff092f54afb73904c924be12 Mon Sep 17 00:00:00 2001
From: Cheng <zcbenz@gmail.com>
Date: Thu, 28 Aug 2025 09:30:08 +0900
Subject: [PATCH] [CUDA] Use ConcurrentContext in concatenate_gpu (#2549)

---
 mlx/backend/cuda/slicing.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlx/backend/cuda/slicing.cpp b/mlx/backend/cuda/slicing.cpp
index 18cc14bbd..93241936b 100644
--- a/mlx/backend/cuda/slicing.cpp
+++ b/mlx/backend/cuda/slicing.cpp
@@ -30,8 +30,7 @@ void concatenate_gpu(
   flags.row_contiguous = false;
   flags.col_contiguous = false;
   flags.contiguous = false;
-  // TODO: Handle concurrent outputs:
-  // https://github.com/ml-explore/mlx/pull/2145#discussion_r2070753816
+  auto concurrent = cu::get_command_encoder(s).concurrent_context();
   for (int i = 0; i < inputs.size(); i++) {
     array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
     size_t data_offset = strides[axis] * sizes[i];