diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt
index 52fecee858..abaca309b2 100644
--- a/mlx/backend/cuda/CMakeLists.txt
+++ b/mlx/backend/cuda/CMakeLists.txt
@@ -174,7 +174,7 @@ endif()
 # Use fixed version of CCCL.
 FetchContent_Declare(
   cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
+  URL "https://github.com/NVIDIA/cccl/releases/download/v3.1.3/cccl-v3.1.3.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
 
diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu
index d164c925b1..c6f82f80a4 100644
--- a/mlx/backend/cuda/sort.cu
+++ b/mlx/backend/cuda/sort.cu
@@ -8,6 +8,7 @@
 
 #include <nvtx3/nvtx3.hpp>
 #include <thrust/device_ptr.h>
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/transform.h>
 #include <cub/device/device_segmented_radix_sort.cuh>