mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-05 02:31:33 +08:00
6f07fa81d7
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com> KVCacheManagerV2 is a new python-based implementation of the KV cache manager, featuring cleaner API, better abstraction and better code quality without the accumulated legacy.
78 lines
2.4 KiB
CMake
Executable File
78 lines
2.4 KiB
CMake
Executable File
set(TRTLLM_PYBIND_MODULE bindings)
|
|
set(TRTLLM_PYBIND_MODULE
|
|
${TRTLLM_PYBIND_MODULE}
|
|
PARENT_SCOPE)
|
|
|
|
set(SRCS
|
|
batch_manager/algorithms.cpp
|
|
batch_manager/bindings.cpp
|
|
batch_manager/buffers.cpp
|
|
batch_manager/cacheTransceiver.cpp
|
|
batch_manager/kvCacheConnector.cpp
|
|
batch_manager/kvCacheManager.cpp
|
|
batch_manager/kvCacheManagerV2Utils.cpp
|
|
batch_manager/llmRequest.cpp
|
|
executor/bindings.cpp
|
|
executor/executor.cpp
|
|
executor/executorConfig.cpp
|
|
executor/request.cpp
|
|
process_group/bindings.cpp
|
|
runtime/bindings.cpp
|
|
runtime/hostfunc.cpp
|
|
common/tllmExceptions.cpp
|
|
testing/modelSpecBinding.cpp
|
|
runtime/moeBindings.cpp
|
|
testing/modelSpecBinding.cpp
|
|
userbuffers/bindings.cpp
|
|
../runtime/ipcNvlsMemory.cu
|
|
thop/bindings.cpp
|
|
bindings.cpp)
|
|
|
|
include_directories(${PROJECT_SOURCE_DIR}/include)
|
|
|
|
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")
|
|
|
|
if(ENABLE_NVSHMEM)
|
|
target_link_libraries(${TRTLLM_PYBIND_MODULE} PUBLIC nvshmem::nvshmem_host
|
|
nvshmem::nvshmem_device)
|
|
endif()
|
|
|
|
target_link_libraries(
|
|
${TRTLLM_PYBIND_MODULE}
|
|
PUBLIC ${SHARED_TARGET}
|
|
${Python3_LIBRARIES}
|
|
${TORCH_LIBRARIES}
|
|
torch_python
|
|
${CUDA_DRV_LIB}
|
|
${CUDA_NVML_LIB}
|
|
th_common
|
|
pg_utils)
|
|
target_compile_definitions(
|
|
${TRTLLM_PYBIND_MODULE} PUBLIC TRTLLM_PYBIND_MODULE=${TRTLLM_PYBIND_MODULE}
|
|
PYBIND11_DETAILED_ERROR_MESSAGES=1)
|
|
|
|
if(NOT WIN32)
|
|
set(TRTLLM_PYBIND_MODULE_RPATH_LIST
|
|
"$ORIGIN/libs" # TRTLLM libraries
|
|
"$ORIGIN/../../.." # Shared libraries under $PREFIX/lib
|
|
"$ORIGIN/../nvidia/nccl/lib" # NCCL libraries
|
|
)
|
|
set_target_properties(
|
|
${TRTLLM_PYBIND_MODULE} PROPERTIES BUILD_RPATH
|
|
"${TRTLLM_PYBIND_MODULE_RPATH_LIST}")
|
|
set_target_properties(
|
|
${TRTLLM_PYBIND_MODULE} PROPERTIES LINK_FLAGS
|
|
"${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}")
|
|
endif()
|
|
|
|
# Build transfer_agent_binding when building bindings (if NIXL is enabled)
|
|
if(TARGET ${TRANSFER_AGENT_BINDING_TARGET})
|
|
add_dependencies(${TRTLLM_PYBIND_MODULE} ${TRANSFER_AGENT_BINDING_TARGET})
|
|
endif()
|