diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6db40118c491..35c368a137f1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -158,10 +158,6 @@ include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files
 add_subdirectory(include)
 add_subdirectory(lib)
 
-# find_package(PythonLibs REQUIRED)
-set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
-set(TRITON_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
-
 # TODO: Figure out which target is sufficient to fix errors; triton is
 # apparently not enough. Currently set linking libstdc++fs for all targets
 # to support some old version GCC compilers like 8.3.0.
@@ -178,22 +174,9 @@ if(TRITON_BUILD_PYTHON_MODULE)
   set(PYTHON_SRC_PATH ${CMAKE_CURRENT_SOURCE_DIR}/python/src)
   include_directories(${PYTHON_SRC_PATH})
 
-  if(PYTHON_INCLUDE_DIRS)
-    # We have PYTHON_INCLUDE_DIRS set--this is what we expect when building
-    # using pip install.
-    include_directories(${PYTHON_INCLUDE_DIRS})
-    include_directories(${PYBIND11_INCLUDE_DIR})
-  else()
-    # Otherwise, we might be building from top CMakeLists.txt directly.
-    # Try to find Python and pybind11 packages.
-    find_package(Python3 REQUIRED COMPONENTS Development Interpreter)
-    find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
-    include_directories(${Python3_INCLUDE_DIRS})
-    include_directories(${pybind11_INCLUDE_DIR})
-    link_directories(${Python3_LIBRARY_DIRS})
-    link_libraries(${Python3_LIBRARIES})
-    add_link_options(${Python3_LINK_OPTIONS})
-  endif()
+  # Python Interpreter is used to run lit tests
+  find_package(Python3 REQUIRED COMPONENTS Development Interpreter)
+  find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
 
   if (DEFINED TRITON_PLUGIN_DIRS)
     foreach(PLUGIN_DIR ${TRITON_PLUGIN_DIRS})
@@ -259,6 +242,9 @@ if(TRITON_BUILD_PYTHON_MODULE)
     LLVMAMDGPUCodeGen
     LLVMAMDGPUAsmParser
 
+    Python3::Module
+    pybind11::headers
+
   )
   if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64" OR # Linux arm64
      CMAKE_SYSTEM_PROCESSOR MATCHES "arm64" OR # macOS arm64
@@ -299,7 +285,7 @@ if(TRITON_BUILD_PYTHON_MODULE)
                   ${PYTHON_SRC_PATH}/llvm.cc)
 
   # Link triton with its dependencies
-  target_link_libraries(triton PUBLIC ${TRITON_LIBRARIES})
+  target_link_libraries(triton PRIVATE ${TRITON_LIBRARIES})
   if(WIN32)
     target_link_libraries(triton PRIVATE ${CMAKE_DL_LIBS})
     set_target_properties(triton PROPERTIES SUFFIX ".pyd")
@@ -322,7 +308,7 @@ if(TRITON_BUILD_PYTHON_MODULE AND NOT WIN32)
     set(PYTHON_LDFLAGS "-undefined dynamic_lookup")
   endif()
 
-  target_link_libraries(triton PRIVATE ${PYTHON_LDFLAGS})
+  target_link_options(triton PRIVATE ${PYTHON_LDFLAGS})
 endif()
 
 if(NOT TRITON_BUILD_PYTHON_MODULE)
diff --git a/python/setup.py b/python/setup.py
index 1fbafb49f26d..698e69e2a6a1 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -400,7 +400,7 @@ def get_pybind11_cmake_args(self):
             pybind11_include_dir = os.path.join(pybind11_sys_path, "include")
         else:
             pybind11_include_dir = pybind11.get_include()
-        return [f"-DPYBIND11_INCLUDE_DIR={pybind11_include_dir}"]
+        return [f"-Dpybind11_INCLUDE_DIR='{pybind11_include_dir}'", f"-Dpybind11_DIR='{pybind11.get_cmake_dir()}'"]
 
     def get_proton_cmake_args(self):
         cmake_args = get_thirdparty_packages([get_json_package_info()])
@@ -438,7 +438,7 @@ def build_extension(self, ext):
             "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON", "-DLLVM_ENABLE_WERROR=ON",
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir, "-DTRITON_BUILD_TUTORIALS=OFF",
             "-DTRITON_BUILD_PYTHON_MODULE=ON", "-DPython3_EXECUTABLE:FILEPATH=" + sys.executable,
-            "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", "-DPYTHON_INCLUDE_DIRS=" + python_include_dir,
+            "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", "-DPython3_INCLUDE_DIR=" + python_include_dir,
             "-DTRITON_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends if not b.is_external]),
             "-DTRITON_PLUGIN_DIRS=" + ';'.join([b.src_dir for b in backends if b.is_external])
         ]
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in
index 4053a8e7df56..fd1fb486dd3c 100644
--- a/test/lit.site.cfg.py.in
+++ b/test/lit.site.cfg.py.in
@@ -2,7 +2,7 @@
 
 import sys
 
-config.triton_obj_root = "@TRITON_BINARY_DIR@"
+config.triton_obj_root = "@triton_BINARY_DIR@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
 config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"
@@ -20,4 +20,4 @@ import lit.llvm
 lit.llvm.initialize(lit_config, config)
 
 # Let the main config do the real work
-lit_config.load_config(config, "@TRITON_SOURCE_DIR@/test/lit.cfg.py")
+lit_config.load_config(config, "@triton_SOURCE_DIR@/test/lit.cfg.py")
diff --git a/third_party/amd/CMakeLists.txt b/third_party/amd/CMakeLists.txt
index 8228c3d39111..a09bab8e1c4c 100644
--- a/third_party/amd/CMakeLists.txt
+++ b/third_party/amd/CMakeLists.txt
@@ -4,6 +4,7 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
   add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms TritonAMDGPUDialectToLLVM)
+  target_link_libraries(TritonAMD PRIVATE Python3::Module pybind11::headers)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)
diff --git a/third_party/nvidia/CMakeLists.txt b/third_party/nvidia/CMakeLists.txt
index 75f98fa8f73a..bab189bcbdd0 100644
--- a/third_party/nvidia/CMakeLists.txt
+++ b/third_party/nvidia/CMakeLists.txt
@@ -4,6 +4,7 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
   add_triton_plugin(TritonNVIDIA ${CMAKE_CURRENT_SOURCE_DIR}/triton_nvidia.cc LINK_LIBS TritonNVIDIAGPUToLLVM NVGPUToLLVM)
+  target_link_libraries(TritonNVIDIA PRIVATE Python3::Module pybind11::headers)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)
diff --git a/third_party/proton/CMakeLists.txt b/third_party/proton/CMakeLists.txt
index e2d9152c9626..ba967c5695a4 100644
--- a/third_party/proton/CMakeLists.txt
+++ b/third_party/proton/CMakeLists.txt
@@ -19,19 +19,8 @@ include_directories(${JSON_INCLUDE_DIR})
 include_directories(${PROTON_SRC_DIR}/include)
 include_directories(${PROTON_EXTERN_DIR})
 
-if(PYTHON_INCLUDE_DIRS)
-  # We have PYTHON_INCLUDE_DIRS set--this is what we expect when building
-  # using pip install.
-  include_directories(${PYTHON_INCLUDE_DIRS})
-  include_directories(${PYBIND11_INCLUDE_DIR})
-else()
-  # Otherwise, we might be building from top CMakeLists.txt directly.
-  # Try to find Python and pybind11 packages.
-  find_package(Python3 REQUIRED Interpreter Development)
-  find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
-  include_directories(${Python3_INCLUDE_DIRS})
-  include_directories(${pybind11_INCLUDE_DIR})
-endif()
+find_package(Python3 REQUIRED Interpreter Development)
+find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
 
 # Check if the platform is MacOS
 if(APPLE)
@@ -49,4 +38,5 @@ include_directories(${CUPTI_INCLUDE_DIR})
 include_directories(SYSTEM ${ROCTRACER_INCLUDE_DIR})
 target_compile_definitions(proton PRIVATE __HIP_PLATFORM_AMD__)
 
-target_link_libraries(proton PRIVATE ${Python_LIBRARIES} ${PROTON_PYTHON_LDFLAGS})
+target_link_libraries(proton PRIVATE Python3::Module pybind11::headers)
+target_link_options(proton PRIVATE ${PROTON_PYTHON_LDFLAGS})
diff --git a/third_party/proton/dialect/CMakeLists.txt b/third_party/proton/dialect/CMakeLists.txt
index c7b5413a0e15..cfa5938873d9 100644
--- a/third_party/proton/dialect/CMakeLists.txt
+++ b/third_party/proton/dialect/CMakeLists.txt
@@ -3,5 +3,6 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
 add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
-  add_triton_plugin(TritonProton ${CMAKE_CURRENT_SOURCE_DIR}/triton_proton.cc LINK_LIBS ProtonIR)
+  add_triton_plugin(TritonProton ${CMAKE_CURRENT_SOURCE_DIR}/triton_proton.cc)
+  target_link_libraries(TritonProton PRIVATE ProtonIR Python3::Module pybind11::headers)
 endif()