diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt
index 7986c09d..fbf60981 100644
--- a/mlx/backend/cuda/CMakeLists.txt
+++ b/mlx/backend/cuda/CMakeLists.txt
@@ -74,8 +74,6 @@ else()
     mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm_batched_12_0.cpp)
 endif()
 
-target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
-
 # Embed kernel sources in binary for JIT compilation.
 file(
   GLOB MLX_JIT_SOURCES
@@ -94,6 +92,10 @@ add_custom_target(cuda_jit_sources DEPENDS gen/cuda_jit_sources.h)
 add_dependencies(mlx cuda_jit_sources)
 target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/gen")
 
+# ------------------------ Compilation configs ------------------------
+
+target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
+
 # Enable defining device lambda functions.
 target_compile_options(mlx
                        PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
@@ -116,6 +118,10 @@ endif()
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
 
+# Suppress nvcc warnings on MLX headers.
+target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
+                                   --diag_suppress=997>)
+
 # Use stronger binaries compression. This feature was introduced in CUDA 12.8
 # and requires drivers released after CUDA 12.4.
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
@@ -143,13 +149,27 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
+# ------------------------ Dependencies ------------------------
+
 # Use fixed version of CCCL.
 FetchContent_Declare(
   cccl
   URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
-set_target_properties(mlx PROPERTIES CCCL_DIR "${cccl_SOURCE_DIR}/include")
+
+# Install CCCL headers for JIT.
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+
+# The binary of C++ tests will not be installed so it can not find the CCCL
+# headers, and we have to hard-code the path.
+if(MLX_BUILD_TESTS)
+  target_compile_definitions(mlx
+                             PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include")
+endif()
 
 # Use fixed version of NVTX.
 FetchContent_Declare(
@@ -187,10 +207,3 @@ target_link_libraries(mlx PRIVATE cudnn_frontend)
 # Link with the actual cuDNN libraries.
 include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
-
-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)
-# Install CCCL headers for JIT.
-install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
-        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp
index 69cf6ff6..cdd04f2f 100644
--- a/mlx/backend/cuda/jit_module.cpp
+++ b/mlx/backend/cuda/jit_module.cpp
@@ -27,55 +27,47 @@ void check_nvrtc_error(const char* name, nvrtcResult err) {
   }
 }
 
-// Return the location of the CUDA toolkit.
-const std::string& cuda_home() {
-  static std::string home = []() -> std::string {
-    const char* home = std::getenv("CUDA_HOME");
-    if (home) {
-      return home;
-    }
-    home = std::getenv("CUDA_PATH");
-    if (home) {
-      return home;
-    }
-#if defined(__linux__)
-    home = "/usr/local/cuda";
-    if (std::filesystem::exists(home)) {
-      return home;
-    }
-#endif
-    throw std::runtime_error(
-        "Environment variable CUDA_HOME or CUDA_PATH is not set.");
-  }();
-  return home;
-}
-
-// Return the location of CCCL headers shipped with the distribution.
-const std::string& cccl_dir() {
-  static std::string dir = []() {
-    std::filesystem::path path;
+// Return the --include-path args used for invoking NVRTC.
+const std::vector<std::string>& include_path_args() {
+  static std::vector<std::string> cached_args = []() {
+    std::vector<std::string> args;
+    // Add path to bundled CCCL headers.
+    auto root_dir = current_binary_dir().parent_path();
+    auto path = root_dir / "include" / "cccl";
 #if defined(MLX_CCCL_DIR)
-    // First search the install dir if defined.
-    path = MLX_CCCL_DIR;
-    if (std::filesystem::exists(path)) {
-      return path.string();
+    if (!std::filesystem::exists(path)) {
+      path = MLX_CCCL_DIR;
     }
 #endif
-    // Then search dynamically from the dir of libmlx.so file.
-    path = current_binary_dir().parent_path() / "include" / "cccl";
     if (std::filesystem::exists(path)) {
-      return path.string();
+      args.push_back(fmt::format("--include-path={}", path.string()));
     }
-    // Finally check the environment variable.
-    if (const char* env = std::getenv("MLX_CCCL_DIR"); env) {
-      path = env;
-      if (!path.empty() && std::filesystem::exists(path)) {
-        return path.string();
+    // Add path to CUDA runtime headers, try local-installed python package
+    // first and then system-installed headers.
+    path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include";
+    if (std::filesystem::exists(path)) {
+      args.push_back(fmt::format("--include-path={}", path.string()));
+    } else {
+      const char* home = std::getenv("CUDA_HOME");
+      if (!home) {
+        home = std::getenv("CUDA_PATH");
+      }
+#if defined(__linux__)
+      if (!home) {
+        home = "/usr/local/cuda";
+      }
+#endif
+      if (home && std::filesystem::exists(home)) {
+        args.push_back(fmt::format("--include-path={}/include", home));
+      } else {
+        throw std::runtime_error(
+            "Can not find locations of CUDA headers, please set environment "
+            "variable CUDA_HOME or CUDA_PATH.");
       }
     }
-    return std::string();
+    return args;
   }();
-  return dir;
+  return cached_args;
 }
 
 // Get the cache directory for storing compiled results.
@@ -288,14 +280,9 @@ void compile(
       device.compute_capability_minor(),
       arch_tag);
   args.push_back(compute.c_str());
-  std::string cccl_include = cccl_dir();
-  if (!cccl_include.empty()) {
-    cccl_include = fmt::format("--include-path={}", cccl_include);
-    args.push_back(cccl_include.c_str());
+  for (const auto& include : include_path_args()) {
+    args.push_back(include.c_str());
   }
-  std::string cuda_include =
-      fmt::format("--include-path={}/include", cuda_home());
-  args.push_back(cuda_include.c_str());
   nvrtcResult compile_result =
       nvrtcCompileProgram(prog, args.size(), args.data());
   if (compile_result != NVRTC_SUCCESS) {
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 74cf74ea..eafa9889 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -36,14 +36,6 @@ target_sources(
           linalg_tests.cpp
           ${METAL_TEST_SOURCES})
 
-if(MLX_BUILD_CUDA)
-  # C++ tests are always built from source, so we have to specify where to find
-  # CCCL headers for JIT as they are not installed in system.
-  get_target_property(MLX_CCCL_DIR mlx CCCL_DIR)
-  target_compile_definitions(mlx PRIVATE MLX_CCCL_DIR="${MLX_CCCL_DIR}")
-  message(STATUS MLX_CCCL_DIR="${MLX_CCCL_DIR}")
-endif()
-
 target_link_libraries(tests PRIVATE mlx doctest)
 doctest_discover_tests(tests)
 add_test(NAME tests COMMAND tests)