From f3aca1ee30f7df3d9a3a33a2a1ba136b1183c77f Mon Sep 17 00:00:00 2001 From: Yang Chen Date: Tue, 1 Apr 2025 06:09:40 -0700 Subject: [PATCH] setup correct nvcc version with CUDA_HOME (#15725) Signed-off-by: Yang Chen --- setup.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 3a92d5a23a1f2..cf2acb20d9cf0 100755 --- a/setup.py +++ b/setup.py @@ -201,6 +201,9 @@ class cmake_build_ext(build_ext): else: # Default build tool to whatever cmake picks. build_tool = [] + # Make sure we use the nvcc from CUDA_HOME + if _is_cuda(): + cmake_args += [f'-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc'] subprocess.check_call( ['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args], cwd=self.build_temp) @@ -639,11 +642,10 @@ if _is_hip(): if _is_cuda(): ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) - if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.0"): - # FA3 requires CUDA 12.0 or later + if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"): + # FA3 requires CUDA 12.3 or later ext_modules.append( CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C")) - if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"): # Optional since this doesn't get built (produce an .so file) when # not targeting a hopper system ext_modules.append(