TensorRT-LLM/cpp/tensorrt_llm/pybind/CMakeLists.txt at main · TheCodeWrangler/TensorRT-LLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
set(TRTLLM_PYBIND_MODULE bindings)
set(TRTLLM_PYBIND_MODULE
    ${TRTLLM_PYBIND_MODULE}
    PARENT_SCOPE)

if(NOT BUILD_PYT)
  message(
    FATAL_ERROR
      "Python bindings for C++ runtime require PyTorch. Please enable BUILD_PYT"
  )
endif()

execute_process(
  COMMAND ${Python3_EXECUTABLE} "-c"
          "import pybind11 as pb11; print(pb11.get_cmake_dir(),end='');"
  RESULT_VARIABLE PYBIND_CMAKE_DIR_RET
  OUTPUT_VARIABLE PYBIND_CMAKE_DIR)

if(PYBIND_CMAKE_DIR_RET MATCHES 0)
  list(APPEND CMAKE_PREFIX_PATH "${PYBIND_CMAKE_DIR}")
else()
  message(ERROR "pybind11 CMake directory not found.")
endif()

find_package(pybind11 REQUIRED)

set(SRCS
    bindings.cpp
    batch_manager/gptManager.cpp
    batch_manager/llmRequest.cpp
    batch_manager/inferenceRequest.cpp
    batch_manager/namedTensor.cpp
    executor/bindings.cpp
    executor/executor.cpp
    runtime/generationInput.cpp
    runtime/generationOutput.cpp)

pybind11_add_module(${TRTLLM_PYBIND_MODULE} ${SRCS})

set_property(TARGET ${TRTLLM_PYBIND_MODULE} PROPERTY POSITION_INDEPENDENT_CODE
                                                     ON)

target_link_directories(${TRTLLM_PYBIND_MODULE} PUBLIC
                        "${TORCH_INSTALL_PREFIX}/lib")
target_link_libraries(
  ${TRTLLM_PYBIND_MODULE}
  PUBLIC ${SHARED_TARGET} ${Python3_LIBRARIES} ${TORCH_LIBRARIES} torch_python
         ${UNDEFINED_FLAG})
target_compile_definitions(${TRTLLM_PYBIND_MODULE}
                           PUBLIC TRTLLM_PYBIND_MODULE=${TRTLLM_PYBIND_MODULE})
target_compile_definitions(${TRTLLM_PYBIND_MODULE}
                           PUBLIC PYBIND11_DETAILED_ERROR_MESSAGES=1)

if(NOT WIN32)
  set_target_properties(
    ${TRTLLM_PYBIND_MODULE}
    PROPERTIES LINK_FLAGS
               "-Wl,-rpath,'$ORIGIN/libs' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}")
endif()