File tree 2 files changed +15
-5
lines changed 2 files changed +15
-5
lines changed Original file line number Diff line number Diff line change @@ -70,10 +70,11 @@ find_package(Torch REQUIRED)
70
70
# config is used for standalone C++ binaries that link against torch).
71
71
# The `libtorch_python.so` library defines some of the glue code between
72
72
# torch/python via pybind and is required by VLLM extensions for this
73
- # reason. So, add it by manually using `append_torchlib_if_found` from
74
- # torch's cmake setup .
73
+ # reason. So, add it by manually with `find_library` using torch's
74
+ # installed library path .
75
75
#
76
- append_torchlib_if_found(torch_python)
76
+ find_library (torch_python_LIBRARY torch_python PATHS
77
+ "${TORCH_INSTALL_PREFIX} /lib" )
77
78
78
79
#
79
80
# Set up GPU language and check the torch version and warn if it isn't
Original file line number Diff line number Diff line change @@ -281,7 +281,7 @@ endmacro()
281
281
# not provided.
282
282
# COMPILE_FLAGS <flags> - Extra compiler flags passed to NVCC/hip.
283
283
# INCLUDE_DIRECTORIES <dirs> - Extra include directories.
284
- # LINK_LIBRARIES <libraries> - Extra link libraries.
284
+ # LIBRARIES <libraries> - Extra link libraries.
285
285
# WITH_SOABI - Generate library with python SOABI suffix name.
286
286
#
287
287
# Note: optimization level/debug info is set via cmake build type.
@@ -327,8 +327,17 @@ function (define_gpu_extension_target GPU_MOD_NAME)
327
327
target_include_directories (${GPU_MOD_NAME} PRIVATE csrc
328
328
${GPU_INCLUDE_DIRECTORIES} )
329
329
330
- target_link_libraries (${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES }
330
+ target_link_libraries (${GPU_MOD_NAME} PRIVATE torch ${torch_python_LIBRARY }
331
331
${GPU_LIBRARIES} )
332
332
333
+ # Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of
334
+ # dependencies that are not necessary and may not be installed.
335
+ if (GPU_LANGUAGE STREQUAL "CUDA" )
336
+ target_link_libraries (${GPU_MOD_NAME} PRIVATE ${CUDA_CUDA_LIB}
337
+ ${CUDA_LIBRARIES} )
338
+ else ()
339
+ target_link_libraries (${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES} )
340
+ endif ()
341
+
333
342
install (TARGETS ${GPU_MOD_NAME} LIBRARY DESTINATION ${GPU_DESTINATION} )
334
343
endfunction ()
You can’t perform that action at this time.
0 commit comments