remove thrust (#3067)

.github/actions/build-linux/action.yml

@@ -9,6 +9,7 @@ inputs:
 runs:
   using: "composite"
   steps:
+
     - name: Install Python package
       id: python_build
       shell: sh

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>)

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

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;

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 {
 

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 {