mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-05 18:51:38 +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.
75 lines
2.3 KiB
CMake
Executable File
75 lines
2.3 KiB
CMake
Executable File
set(TRTLLM_NB_MODULE bindings)
|
|
set(TRTLLM_NB_MODULE
|
|
${TRTLLM_NB_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
|
|
common/tllmExceptions.cpp
|
|
executor/bindings.cpp
|
|
executor/executor.cpp
|
|
executor/executorConfig.cpp
|
|
executor/request.cpp
|
|
process_group/bindings.cpp
|
|
runtime/bindings.cpp
|
|
runtime/hostfunc.cpp
|
|
runtime/moeBindings.cpp
|
|
testing/modelSpecBinding.cpp
|
|
userbuffers/bindings.cpp
|
|
thop/bindings.cpp
|
|
../runtime/ipcNvlsMemory.cu
|
|
bindings.cpp)
|
|
|
|
include_directories(${PROJECT_SOURCE_DIR}/include)
|
|
|
|
nanobind_add_module(${TRTLLM_NB_MODULE} ${SRCS})
|
|
|
|
set_property(TARGET ${TRTLLM_NB_MODULE} PROPERTY POSITION_INDEPENDENT_CODE ON)
|
|
|
|
target_link_directories(${TRTLLM_NB_MODULE} PUBLIC
|
|
"${TORCH_INSTALL_PREFIX}/lib")
|
|
|
|
if(ENABLE_NVSHMEM)
|
|
target_link_libraries(${TRTLLM_NB_MODULE} PUBLIC nvshmem::nvshmem_host
|
|
nvshmem::nvshmem_device)
|
|
endif()
|
|
|
|
target_link_libraries(
|
|
${TRTLLM_NB_MODULE}
|
|
PUBLIC ${SHARED_TARGET}
|
|
${Python3_LIBRARIES}
|
|
${TORCH_LIBRARIES}
|
|
torch_python
|
|
${CUDA_DRV_LIB}
|
|
${CUDA_NVML_LIB}
|
|
th_common
|
|
pg_utils)
|
|
target_compile_definitions(
|
|
${TRTLLM_NB_MODULE} PUBLIC TRTLLM_NB_MODULE=${TRTLLM_NB_MODULE}
|
|
PYBIND11_DETAILED_ERROR_MESSAGES=1)
|
|
|
|
if(NOT WIN32)
|
|
set(TRTLLM_NB_MODULE_RPATH_LIST
|
|
"$ORIGIN/libs" # TRTLLM libraries
|
|
"$ORIGIN/../../.." # Shared libraries under $PREFIX/lib
|
|
"$ORIGIN/../nvidia/nccl/lib" # NCCL libraries
|
|
)
|
|
set_target_properties(${TRTLLM_NB_MODULE}
|
|
PROPERTIES BUILD_RPATH "${TRTLLM_NB_MODULE_RPATH_LIST}")
|
|
set_target_properties(
|
|
${TRTLLM_NB_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_NB_MODULE} ${TRANSFER_AGENT_BINDING_TARGET})
|
|
endif()
|