From 2f324cc3b200700b422db4811ae3ff8bd5bf48b4 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 27 Jan 2026 08:54:07 -0800
Subject: [PATCH] remove thrust (#3067)

---
 .github/actions/build-linux/action.yml | 1 +
 mlx/backend/cuda/CMakeLists.txt        | 4 ++--
 mlx/backend/cuda/device.h              | 8 --------
 mlx/backend/cuda/quantized/qqmm.cpp    | 1 -
 mlx/backend/cuda/reduce/col_reduce.cu  | 1 +
 mlx/backend/cuda/utils.cpp             | 1 +
 6 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/.github/actions/build-linux/action.yml b/.github/actions/build-linux/action.yml
index c7a74713..166d7b8a 100644
--- a/.github/actions/build-linux/action.yml
+++ b/.github/actions/build-linux/action.yml
@@ -9,6 +9,7 @@ inputs:
 runs:
   using: "composite"
   steps:
+
     - name: Install Python package
       id: python_build
       shell: sh
diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt
index 601fbf74..013b24b2 100644
--- a/mlx/backend/cuda/CMakeLists.txt
+++ b/mlx/backend/cuda/CMakeLists.txt
@@ -257,9 +257,9 @@ target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 FetchContent_Declare(
   cutlass
   GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-  GIT_TAG v4.3.2
+  GIT_TAG v4.3.5
   GIT_SHALLOW TRUE
   SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(cutlass)
 target_include_directories(
-  mlx PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
+  mlx SYSTEM PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h
index 81fd1d22..41d6cf8f 100644
--- a/mlx/backend/cuda/device.h
+++ b/mlx/backend/cuda/device.h
@@ -11,7 +11,6 @@
 #include <cublasLt.h>
 #include <cuda.h>
 #include <cudnn.h>
-#include <thrust/execution_policy.h>
 
 #include <unordered_map>
 
@@ -181,11 +180,4 @@ Device& device(int cuda_device);
 Device& device(mlx::core::Device d);
 CommandEncoder& get_command_encoder(Stream s);
 
-// Return an execution policy that does not sync for result.
-// Note that not all thrust APIs support async policy, confirm before using.
-inline auto thrust_policy(cudaStream_t stream) {
-  // TODO: Connect thrust's custom allocator with mlx's allocator.
-  return thrust::cuda::par_nosync.on(stream);
-}
-
 } // namespace mlx::core::cu
diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp
index 93724b2f..3b5be771 100644
--- a/mlx/backend/cuda/quantized/qqmm.cpp
+++ b/mlx/backend/cuda/quantized/qqmm.cpp
@@ -78,7 +78,6 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     fp_qmv(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder);
     return;
   }
-  std::cout << "RUNNING FULL?" << std::endl;
 
   auto cc = device.compute_capability_major() * 100 +
       device.compute_capability_minor() * 10;
diff --git a/mlx/backend/cuda/reduce/col_reduce.cu b/mlx/backend/cuda/reduce/col_reduce.cu
index 1188faca..e33551d8 100644
--- a/mlx/backend/cuda/reduce/col_reduce.cu
+++ b/mlx/backend/cuda/reduce/col_reduce.cu
@@ -8,6 +8,7 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <cub/block/block_load.cuh>
+#include <cub/cub.cuh>
 
 namespace mlx::core {
 
diff --git a/mlx/backend/cuda/utils.cpp b/mlx/backend/cuda/utils.cpp
index 7640f903..08f0f3d4 100644
--- a/mlx/backend/cuda/utils.cpp
+++ b/mlx/backend/cuda/utils.cpp
@@ -5,6 +5,7 @@
 #include "mlx/dtype_utils.h"
 
 #include <fmt/format.h>
+#include <cuda/cmath>
 #include <vector>
 
 namespace mlx::core {