[CUDA] Add more ways finding CCCL headers in JIT

tests/CMakeLists.txt

@@ -39,6 +39,14 @@ target_sources(
           linalg_tests.cpp
           ${METAL_TEST_SOURCES})
 
+if(MLX_BUILD_CUDA)
+  # Find the CCCL headers in install dir.
+  target_compile_definitions(
+    mlx
+    PRIVATE
+      MLX_CCCL_DIR="${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_INCLUDEDIR}/cccl")
+endif()
+
 target_link_libraries(tests PRIVATE mlx doctest)
 doctest_discover_tests(tests)
 add_test(NAME tests COMMAND tests)