From 1b32e0264888200a0e6187496a816ef597a7f320 Mon Sep 17 00:00:00 2001
From: sasha0552 <admin@sasha0552.org>
Date: Wed, 21 Aug 2024 18:17:48 +0000
Subject: [PATCH] [Bugfix] Pass PYTHONPATH from setup.py to CMake (#7730)

---
 CMakeLists.txt | 2 +-
 setup.py       | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c8d4aaeda9091..217dc70c4b24e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -233,7 +233,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # Generate sources:
   execute_process(
     COMMAND ${CMAKE_COMMAND} -E env 
-    PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:$PYTHONPATH 
+    PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$PYTHONPATH 
       ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py
     RESULT_VARIABLE machete_generation_result
     OUTPUT_VARIABLE machete_generation_output
diff --git a/setup.py b/setup.py
index ef599b613667b..21b0422c0f0bd 100644
--- a/setup.py
+++ b/setup.py
@@ -184,6 +184,10 @@ class cmake_build_ext(build_ext):
         # match.
         cmake_args += ['-DVLLM_PYTHON_EXECUTABLE={}'.format(sys.executable)]
 
+        # Pass the python path to cmake so it can reuse the build dependencies
+        # on subsequent calls to python.
+        cmake_args += ['-DVLLM_PYTHON_PATH={}'.format(":".join(sys.path))]
+
         #
         # Setup parallelism and build tool
         #