A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
setup correct nvcc version with CUDA_HOME (#15725)
Signed-off-by: Yang Chen <yangche@fb.com>
Yang Chen committed
Apr 1, 2025 at 06:09 UTC
f3aca1ee30f7df3d9a3a33a2a1ba136b1183c77f
1 file changed
+5
-3
setup.py
+5
-3
@@ -201,6 +201,9 @@ class cmake_build_ext(build_ext):
201
else:
202
# Default build tool to whatever cmake picks.
203
build_tool = []
204
+ # Make sure we use the nvcc from CUDA_HOME
205
+ if _is_cuda():
206
+ cmake_args += [f'-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc']
207
subprocess.check_call(
208
['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args],
209
cwd=self.build_temp)
@@ -639,11 +642,10 @@ if _is_hip():
642
643
if _is_cuda():
644
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
642
- if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.0"):
643
- # FA3 requires CUDA 12.0 or later
645
+ if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
646
+ # FA3 requires CUDA 12.3 or later
647
ext_modules.append(
648
CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
646
- if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
649
# Optional since this doesn't get built (produce an .so file) when
650
# not targeting a hopper system
651
ext_modules.append(