diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml index 4b8cc1acac63..1619f3e6c9ff 100644 --- a/.devcontainer/docker-compose.yml +++ b/.devcontainer/docker-compose.yml @@ -1,7 +1,7 @@ version: "3.9" services: tensorrt_llm-dev: - image: urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505121727-4049 + image: urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505160532-3934 network_mode: host ipc: host diff --git a/cpp/cmake/modules/FindNIXL.cmake b/cpp/cmake/modules/FindNIXL.cmake new file mode 100644 index 000000000000..288add626aa5 --- /dev/null +++ b/cpp/cmake/modules/FindNIXL.cmake @@ -0,0 +1,73 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. +# + +# Check if NIXL is already found +if(NIXL_FOUND) + # If NIXL is already found, exit the script early + return() +endif() + +find_package(ucx REQUIRED) + +find_path(NIXL_INCLUDE_DIR nixl.h HINTS ${NIXL_ROOT}/include) + +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + set(NIXL_TARGET_ARCH "x86_64-linux-gnu") +elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + set(NIXL_TARGET_ARCH "aarch64-linux-gnu") +endif() + +find_library(NIXL_LIBRARY nixl HINTS ${NIXL_ROOT}/lib/${NIXL_TARGET_ARCH} + ${NIXL_ROOT}/lib64) +find_library(NIXL_BUILD_LIBRARY nixl_build + HINTS ${NIXL_ROOT}/lib/${NIXL_TARGET_ARCH} ${NIXL_ROOT}/lib64) +find_library(SERDES_LIBRARY serdes HINTS ${NIXL_ROOT}/lib/${NIXL_TARGET_ARCH} + ${NIXL_ROOT}/lib64) +find_library( + UCX_BACKEND_LIBRARY plugin_UCX + HINTS ${NIXL_ROOT}/lib/${NIXL_TARGET_ARCH}/plugins ${NIXL_ROOT}/lib64/plugins) +find_library(UCX_UTILS_LIBRARY ucx_utils + HINTS ${NIXL_ROOT}/lib/${NIXL_TARGET_ARCH} ${NIXL_ROOT}/lib64) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args( + NIXL + FOUND_VAR NIXL_FOUND + REQUIRED_VARS NIXL_INCLUDE_DIR NIXL_LIBRARY NIXL_BUILD_LIBRARY SERDES_LIBRARY) + +# Re-attempt to find NIXL after installation +find_package_handle_standard_args( + NIXL + FOUND_VAR NIXL_FOUND + REQUIRED_VARS NIXL_INCLUDE_DIR NIXL_LIBRARY NIXL_BUILD_LIBRARY SERDES_LIBRARY) + +# Set up the NIXL target if found +if(NIXL_FOUND) + if(NOT TARGET NIXL::nixl) + add_library(NIXL::nixl SHARED IMPORTED) + set_target_properties( + NIXL::nixl + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${NIXL_INCLUDE_DIR} + IMPORTED_LOCATION ${NIXL_LIBRARY} ${NIXL_BUILD_LIBRARY} + ${SERDES_LIBRARY}) + endif() +else() + message(STATUS "NIXL_LIBRARY: ${NIXL_LIBRARY}") + message(STATUS "NIXL_BUILD_LIBRARY: ${NIXL_BUILD_LIBRARY}") + message(STATUS "SERDES_LIBRARY: ${SERDES_LIBRARY}") + message(FATAL_ERROR "NIXL not found after installation attempt.") +endif() diff --git a/cpp/include/tensorrt_llm/common/assert.h b/cpp/include/tensorrt_llm/common/assert.h index 768bb3e212ba..e7e24bf549e9 100644 --- a/cpp/include/tensorrt_llm/common/assert.h +++ b/cpp/include/tensorrt_llm/common/assert.h @@ -23,6 +23,11 @@ namespace tensorrt_llm::common { +[[noreturn]] inline void throwRuntimeError(char const* const file, int const line, char const* info) +{ + throw TllmException(file, line, fmtstr("[TensorRT-LLM][ERROR] Assertion failed: %s", info).c_str()); +} + [[noreturn]] inline void throwRuntimeError(char const* const file, int const line, std::string const& info = "") { throw TllmException(file, line, fmtstr("[TensorRT-LLM][ERROR] Assertion failed: %s", info.c_str()).c_str()); @@ -57,7 +62,7 @@ class DebugConfig TLLM_LIKELY(static_cast(val)) \ ? ((void) 0) \ : tensorrt_llm::common::throwRuntimeError( \ - __FILE__, __LINE__, tensorrt_llm::common::fmtstr(info, ##__VA_ARGS__)); \ + __FILE__, __LINE__, tensorrt_llm::common::fmtstr(info, ##__VA_ARGS__).c_str()); \ } while (0) #define TLLM_CHECK_DEBUG(val) \ @@ -78,7 +83,7 @@ class DebugConfig TLLM_LIKELY(static_cast(val)) \ ? ((void) 0) \ : tensorrt_llm::common::throwRuntimeError( \ - __FILE__, __LINE__, tensorrt_llm::common::fmtstr(info, ##__VA_ARGS__)); \ + __FILE__, __LINE__, tensorrt_llm::common::fmtstr(info, ##__VA_ARGS__).c_str()); \ } \ } while (0) diff --git a/cpp/include/tensorrt_llm/executor/cacheCommunicator.h b/cpp/include/tensorrt_llm/executor/cacheCommunicator.h index 7290592b7530..345650d84d0a 100644 --- a/cpp/include/tensorrt_llm/executor/cacheCommunicator.h +++ b/cpp/include/tensorrt_llm/executor/cacheCommunicator.h @@ -16,21 +16,14 @@ #pragma once -#include -#include - -#include "tensorrt_llm/executor/dataTransceiverState.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -namespace tensorrt_llm::batch_manager -{ -class RequestInfo; -class UcxEndpoint; -} // namespace tensorrt_llm::batch_manager +#include +#include namespace tensorrt_llm::executor::kv_cache { +class CommState; + struct DataContext { public: diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h new file mode 100644 index 000000000000..c50605334116 --- /dev/null +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -0,0 +1,259 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache +{ + +enum class MemoryType : uint8_t +{ + kDRAM, + kVRAM, + kBLK, + kOBJ, + kFILE +}; + +class MemoryDesc +{ +public: + MemoryDesc(std::vector const& vec, uint32_t deviceId = 0) + : mAddr{reinterpret_cast(vec.data())} + , mLen{vec.size()} + , mDeviceId{deviceId} + { + } + + MemoryDesc(void* addr, size_t len, uint32_t deviceId) + : mAddr{reinterpret_cast(addr)} + , mLen{len} + , mDeviceId{deviceId} + { + } + + MemoryDesc(uintptr_t addr, size_t len, uint32_t deviceId) + : mAddr{addr} + , mLen{len} + , mDeviceId{deviceId} + { + } + + [[nodiscard]] uintptr_t getAddr() const noexcept + { + return mAddr; + } + + [[nodiscard]] size_t getLen() const noexcept + { + return mLen; + } + + [[nodiscard]] uint32_t getDeviceId() const noexcept + { + return mDeviceId; + } + +private: + uintptr_t mAddr; + size_t mLen; + uint32_t mDeviceId; +}; + +class MemoryDescs +{ +public: + MemoryDescs(MemoryType type, std::vector descs) + : mType{type} + , mDescs{std::move(descs)} + { + } + + [[nodiscard]] MemoryType getType() const noexcept + { + return mType; + } + + [[nodiscard]] std::vector const& getDescs() const noexcept + { + return mDescs; + } + +private: + MemoryType mType; + std::vector mDescs; +}; + +using TransferDescs = MemoryDescs; +using RegisterDescs = MemoryDescs; + +class AgentDesc final +{ +public: + AgentDesc(std::vector backendAgentDesc) + : mBackendAgentDesc{backendAgentDesc} + { + } + + [[nodiscard]] std::vector const& getBackendAgentDesc() const noexcept + { + return mBackendAgentDesc; + } + +private: + std::vector mBackendAgentDesc; +}; + +enum class TransferOp : uint8_t +{ + kREAD, + kWRITE, +}; + +class TransferRequest +{ +public: + TransferRequest(TransferOp op, TransferDescs srcDescs, TransferDescs dstDescs, char const* remoteName) + : mOp{op} + , mSrcDescs{std::move(srcDescs)} + , mDstDescs{std::move(dstDescs)} + , mRemoteName{remoteName} + { + } + + [[nodiscard]] TransferOp getOp() const noexcept + { + return mOp; + } + + [[nodiscard]] TransferDescs const& getSrcDescs() const noexcept + { + return mSrcDescs; + } + + [[nodiscard]] TransferDescs const& getDstDescs() const noexcept + { + return mDstDescs; + } + + [[nodiscard]] char const* getRemoteName() const noexcept + { + return mRemoteName.c_str(); + } + +private: + TransferOp mOp; + TransferDescs mSrcDescs; + TransferDescs mDstDescs; + std::string mRemoteName; +}; + +class TransferStatus +{ +public: + virtual ~TransferStatus() = default; + [[nodiscard]] virtual bool isCompleted() const = 0; + virtual void wait() const = 0; +}; + +class AgentRegistrar +{ +public: + ~AgentRegistrar() = default; + + [[nodiscard]] virtual AgentDesc const* getAgentDesc(char const* agentName) const = 0; + + virtual void addAgentDesc(char const* agentName, AgentDesc desc) = 0; + + virtual void removeAgentDesc(char const* agentName) = 0; +}; + +struct BaseAgentConfig +{ + char const* mName; + bool useProgThread; +}; + +class BaseTransferAgent +{ +public: + virtual ~BaseTransferAgent() = default; + + virtual void registerMemory(RegisterDescs const& descs) = 0; + + virtual void deregisterMemory(RegisterDescs const& descs) = 0; + + virtual void loadRemoteAgent(char const* name) = 0; + + virtual void invalidateRemoteAgent(char const* name) = 0; + + [[nodiscard]] virtual std::unique_ptr submitTransferRequests(TransferRequest const& request) = 0; + + // TODO: Add `notifySyncInfo` and `getMatchedSyncInfo` interfaces. +}; + +class DynLibLoader final +{ +public: + [[nodiscard]] static DynLibLoader& getInstance(); + + [[nodiscard]] void* getHandle(char const* name); + + template + [[nodiscard]] FunctionT getFunctionPointer(char const* libName, char const* funcName) + { + void* handle = getHandle(libName); + void* funcPtr = dlSym(handle, funcName); + const std::string err = funcName + std::string{" function is not open correctly."}; + TLLM_CHECK_WITH_INFO(funcPtr, "%s", err.c_str()); + return reinterpret_cast(funcPtr); + } + + ~DynLibLoader(); + + DynLibLoader() = default; + DynLibLoader(DynLibLoader const&) = delete; + DynLibLoader& operator=(DynLibLoader const&) = delete; + +private: + [[nodiscard]] static void* dlSym(void* handle, char const* symbol); + + std::mutex mDllMutex; + std::unordered_map mHandlers; +}; + +template +[[nodiscard]] std::unique_ptr makeTransferAgent(char const* const& backend, Args&&... args) +{ + if (backend == std::string{"nixl"}) + { + auto& loader = DynLibLoader::getInstance(); + using CreateNixlFuncType = std::unique_ptr (*)(BaseAgentConfig const*, AgentRegistrar*); + auto* func = loader.getFunctionPointer( + "libtensorrt_llm_nixl_wrapper.so", "createNixlTransferAgent"); + return func(std::forward(args)...); + } + TLLM_THROW("Unknown backend name."); +} + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 0db57ad814fe..49ae14c36d8d 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -148,6 +148,11 @@ add_subdirectory(batch_manager) set(EXECUTOR_TARGET tensorrt_llm_executor_static) set(EXECUTOR_TARGET_ARCH ${TARGET_ARCH}) set(UCX_WRAPPER_TARGET tensorrt_llm_ucx_wrapper) + +if(NIXL_ROOT) + set(NIXL_WRAPPER_TARGET tensorrt_llm_nixl_wrapper) +endif() + add_subdirectory(executor) find_package(Threads REQUIRED) @@ -233,6 +238,11 @@ if(TARGET ${UCX_WRAPPER_TARGET}) add_dependencies(${SHARED_TARGET} ${UCX_WRAPPER_TARGET}) endif() +if(TARGET ${NIXL_WRAPPER_TARGET}) + target_link_libraries(${NIXL_WRAPPER_TARGET} INTERFACE ${SHARED_TARGET}) + add_dependencies(${SHARED_TARGET} ${NIXL_WRAPPER_TARGET}) +endif() + if(NOT WIN32) set_target_properties(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'") diff --git a/cpp/tensorrt_llm/executor/CMakeLists.txt b/cpp/tensorrt_llm/executor/CMakeLists.txt index f6d8b72d39b0..56e474ab8c43 100644 --- a/cpp/tensorrt_llm/executor/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/CMakeLists.txt @@ -16,6 +16,7 @@ set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}) # keep this list sorted alphabetically set(SRCS cache_transmission/mpi_utils/connection.cpp + cache_transmission/transferAgent.cpp cache_transmission/cacheConcatenate.cu contextPhaseParams.cpp debugConfig.cpp @@ -83,3 +84,4 @@ target_compile_definitions(${EXECUTOR_STATIC_TARGET} PUBLIC TOP_LEVEL_DIR="${TOP_LEVEL_DIR}") add_subdirectory(cache_transmission/ucx_utils) +add_subdirectory(cache_transmission/nixl_utils) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h index 6a6472d0f155..aca83131ec21 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/cacheCommunicator.h" +#include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt new file mode 100644 index 000000000000..cff9d767e85f --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. SPDX-License-Identifier: NVIDIA TensorRT +# Source Code License Agreement +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this material and related documentation without an express +# license agreement from NVIDIA CORPORATION or its affiliates is strictly +# prohibited. + +if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + message( + STATUS + "The NIXL backend is temporarily unavailable on the aarch64 platform.") + unset(NIXL_ROOT) +endif() + +if(NIXL_ROOT) + find_package(NIXL REQUIRED) + # Check if all required packages were found + if(NOT NIXL_FOUND) + message( + FATAL_ERROR "NIXL not found. Please install NIXL or set `NIXL_ROOT`.") + endif() + + # Define the NIXL wrapper target + set(NIXL_WRAPPER_TARGET "tensorrt_llm_nixl_wrapper") + + # Add the NIXL wrapper target + add_library(${NIXL_WRAPPER_TARGET} SHARED transferAgent.cpp) + target_compile_options( + ${NIXL_WRAPPER_TARGET} PRIVATE -U_GLIBCXX_USE_CXX11_ABI + -D_GLIBCXX_USE_CXX11_ABI=1 -Wno-error) + + # Add include directories + target_include_directories(${NIXL_WRAPPER_TARGET} PRIVATE NIXL::nixl) + + # Link against all NIXL libraries + target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE NIXL::nixl) + + add_custom_command( + TARGET ${NIXL_WRAPPER_TARGET} + POST_BUILD + COMMAND nm -C $ | grep -E + 'tensorrt_llm.*__cxx11|__cxx11.*tensorrt_llm' && exit 1 || exit 0) +endif() diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/interfaces.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/interfaces.h new file mode 100644 index 000000000000..76a9320b226a --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/interfaces.h @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/executor/transferAgent.h" + +namespace tensorrt_llm::executor::kv_cache +{ + +#if defined(__clang__) +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wreturn-type-c-linkage" +#endif + +extern "C" +{ + [[nodiscard]] std::unique_ptr createNixlTransferAgent( + BaseAgentConfig const* config, AgentRegistrar* registrar); +} + +#if defined(__clang__) +#pragma clang diagnostic pop +#endif + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp new file mode 100644 index 000000000000..faa02355eb25 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -0,0 +1,180 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" + +namespace tensorrt_llm::executor::kv_cache +{ + +[[nodiscard]] nixl_mem_t NixlHelper::convert(MemoryType type) +{ + switch (type) + { + case MemoryType::kDRAM: return DRAM_SEG; + case MemoryType::kVRAM: return VRAM_SEG; + case MemoryType::kBLK: return BLK_SEG; + case MemoryType::kOBJ: return OBJ_SEG; + case MemoryType::kFILE: return FILE_SEG; + default: TLLM_THROW("Unknown MemoryType value"); + } +} + +[[nodiscard]] nixlBasicDesc NixlHelper::convert(MemoryDesc const& desc) +{ + return nixlBasicDesc{desc.getAddr(), desc.getLen(), desc.getDeviceId()}; +} + +[[nodiscard]] nixl_reg_dlist_t NixlHelper::convertRegDlist(RegisterDescs const& descs) +{ + nixl_reg_dlist_t list{convert(descs.getType())}; + for (auto const& desc : descs.getDescs()) + { + list.addDesc(nixlBlobDesc{desc.getAddr(), desc.getLen(), desc.getDeviceId()}); + } + return list; +} + +[[nodiscard]] nixl_xfer_op_t NixlHelper::convert(TransferOp const& op) +{ + switch (op) + { + case TransferOp::kREAD: return NIXL_READ; + case TransferOp::kWRITE: return NIXL_WRITE; + default: TLLM_THROW("Unknown TransferOp value"); + } +} + +[[nodiscard]] nixl_xfer_dlist_t NixlHelper::convertXferDist(TransferDescs const& descs) +{ + nixl_xfer_dlist_t list{convert(descs.getType())}; + for (auto const& desc : descs.getDescs()) + { + list.addDesc(nixlBasicDesc{desc.getAddr(), desc.getLen(), desc.getDeviceId()}); + } + return list; +} + +NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) + : mRawAgent{agent} + , mHandle{handle} +{ + TLLM_CHECK(mRawAgent); + TLLM_CHECK(mHandle); +} + +void NixlTransferStatus::wait() const +{ + while (!isCompleted()) + ; +} + +[[nodiscard]] bool NixlTransferStatus::isCompleted() const +{ + return mRawAgent->getXferStatus(mHandle) == NIXL_SUCCESS; +} + +NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config, AgentRegistrar* registrar) + : mRegistrar{registrar} + , mName{config.mName} +{ + nixl_status_t status; + TLLM_CHECK(mRegistrar); + nixlAgentConfig nixlConfig{config.useProgThread}; + mRawAgent = std::make_unique(config.mName, std::move(nixlConfig)); + + nixl_b_params_t init1; + nixl_mem_list_t mems1; + status = mRawAgent->getPluginParams("UCX", mems1, init1); + TLLM_CHECK(status == NIXL_SUCCESS); + + status = mRawAgent->createBackend("UCX", init1, mRawBackend); + if (status != NIXL_SUCCESS || !mRawBackend) + { + TLLM_THROW("Failed to create NIXL backend"); + } + mExtraParams.backends.push_back(mRawBackend); +} + +void NixlTransferAgent::registerMemory(RegisterDescs const& descs) +{ + nixl_status_t status; + status = mRawAgent->registerMem(NixlHelper::convertRegDlist(descs), &mExtraParams); + TLLM_CHECK(status == NIXL_SUCCESS); + + std::string localMD; + status = mRawAgent->getLocalMD(localMD); + TLLM_CHECK(status == NIXL_SUCCESS); + mRegistrar->addAgentDesc(mName.c_str(), std::vector(localMD.begin(), localMD.end())); +} + +void NixlTransferAgent::deregisterMemory(RegisterDescs const& descs) +{ + nixl_status_t status; + status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(descs), &mExtraParams); + TLLM_CHECK(status == NIXL_SUCCESS); +} + +void NixlTransferAgent::loadRemoteAgent(char const* name) +{ + nixl_status_t status; + auto const* desc = mRegistrar->getAgentDesc(name); + TLLM_CHECK(desc); + std::string remoteName; + auto backendDesc = desc->getBackendAgentDesc(); + status = mRawAgent->loadRemoteMD(std::string(backendDesc.begin(), backendDesc.end()), remoteName); + TLLM_CHECK(status == NIXL_SUCCESS); + TLLM_CHECK_WITH_INFO( + name == remoteName, "loadRemoteAgent gets error agent name: %s != %s", name, remoteName.c_str()); +} + +void NixlTransferAgent::invalidateRemoteAgent(char const* name) +{ + mRawAgent->invalidateRemoteMD(name); +} + +[[nodiscard]] std::unique_ptr NixlTransferAgent::submitTransferRequests(TransferRequest const& request) +{ + nixl_status_t status; + nixlXferReqH* handle; + status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), + NixlHelper::convertXferDist(request.getSrcDescs()), NixlHelper::convertXferDist(request.getDstDescs()), + request.getRemoteName(), handle, &mExtraParams); + TLLM_CHECK(status == NIXL_SUCCESS); + + status = mRawAgent->postXferReq(handle, &mExtraParams); + return std::make_unique(mRawAgent.get(), handle); +} + +#if defined(__clang__) +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wreturn-type-c-linkage" +#endif + +extern "C" +{ + std::unique_ptr createNixlTransferAgent(BaseAgentConfig const* config, AgentRegistrar* registrar) + { + TLLM_CHECK(config); + return std::make_unique(*config, registrar); + } +} + +#if defined(__clang__) +#pragma clang diagnostic pop +#endif + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h new file mode 100644 index 000000000000..4b64ec63dc95 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "nixl.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/interfaces.h" +#include "tensorrt_llm/executor/transferAgent.h" + +namespace tensorrt_llm::executor::kv_cache +{ + +struct NixlHelper +{ + [[nodiscard]] static nixl_mem_t convert(MemoryType type); + [[nodiscard]] static nixlBasicDesc convert(MemoryDesc const& desc); + [[nodiscard]] static nixl_reg_dlist_t convertRegDlist(RegisterDescs const& descs); + [[nodiscard]] static nixl_xfer_op_t convert(TransferOp const& op); + [[nodiscard]] static nixl_xfer_dlist_t convertXferDist(TransferDescs const& descs); +}; + +class NixlTransferStatus final : public TransferStatus +{ +public: + NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle); + + [[nodiscard]] bool isCompleted() const override; + + void wait() const override; + +private: + nixlAgent* mRawAgent{}; + nixlXferReqH* mHandle{}; +}; + +class NixlTransferAgent final : public BaseTransferAgent +{ +public: + NixlTransferAgent(BaseAgentConfig const& config, AgentRegistrar* registrar); + + void registerMemory(RegisterDescs const& descs) override; + + void deregisterMemory(RegisterDescs const& descs) override; + + void loadRemoteAgent(char const* name) override; + + void invalidateRemoteAgent(char const* name) override; + + [[nodiscard]] std::unique_ptr submitTransferRequests(TransferRequest const& request) override; + + [[nodiscard]] nixlAgent* getRawAgent() const noexcept + { + return mRawAgent.get(); + } + + nixl_opt_args_t* getExtraParams() noexcept + { + return &mExtraParams; + } + +private: + std::unique_ptr mRawAgent; + nixlBackendH* mRawBackend{}; + AgentRegistrar* mRegistrar{}; + nixl_opt_args_t mExtraParams; + std::string mName; +}; + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp new file mode 100644 index 000000000000..3e7790cef1fd --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/executor/transferAgent.h" +#include "tensorrt_llm/common/logger.h" + +#include + +namespace tensorrt_llm::executor::kv_cache +{ + +[[nodiscard]] DynLibLoader& DynLibLoader::getInstance() +{ + static DynLibLoader instance; + return instance; +} + +[[nodiscard]] void* DynLibLoader::getHandle(char const* name) +{ + const std::string nameStr{name}; + std::lock_guard lock(mDllMutex); + auto it = mHandlers.find(nameStr); + if (it != mHandlers.end()) + { + return it->second; + } + TLLM_LOG_INFO("dlopen: " + nameStr); + void* handler = dlopen(name, RTLD_LAZY); + TLLM_CHECK_WITH_INFO(handler, "%s can not be loaded correctly: %s", name, dlerror()); + mHandlers[nameStr] = handler; + return handler; +} + +[[nodiscard]] void* DynLibLoader::dlSym(void* handle, char const* symbol) +{ + return dlsym(handle, symbol); +} + +DynLibLoader::~DynLibLoader() +{ + std::lock_guard lock(mDllMutex); + for (auto const& pair : mHandlers) + { + dlclose(pair.second); + } +} + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/connection.h index f2ef4f877439..5bdee4be2106 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/connection.h @@ -21,6 +21,7 @@ #include "ucxx/utils/sockaddr.h" #include "ucxx/utils/ucx.h" #include +#include #if __linux__ #include #include diff --git a/cpp/tests/executor/CMakeLists.txt b/cpp/tests/executor/CMakeLists.txt index 765bbb22fae5..885247a92887 100644 --- a/cpp/tests/executor/CMakeLists.txt +++ b/cpp/tests/executor/CMakeLists.txt @@ -16,3 +16,8 @@ add_gtest(encDecTest encDecTest.cpp) target_link_libraries(encDecTest PRIVATE testingUtils) add_gtest(disaggExecutorTest disaggExecutorTest.cpp) target_link_libraries(disaggExecutorTest PRIVATE testingUtils) + +if(NIXL_ROOT) + add_gtest(transferAgentTest cache_transmission/transferAgentTest.cpp) + target_link_libraries(transferAgentTest PRIVATE tensorrt_llm_nixl_wrapper) +endif() diff --git a/cpp/tests/executor/cache_transmission/transferAgentTest.cpp b/cpp/tests/executor/cache_transmission/transferAgentTest.cpp new file mode 100644 index 000000000000..ada271d4dc81 --- /dev/null +++ b/cpp/tests/executor/cache_transmission/transferAgentTest.cpp @@ -0,0 +1,112 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: NVIDIA TensorRT Source Code License Agreement + * + * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual + * property and proprietary rights in and to this material, related + * documentation and any modifications thereto. Any use, reproduction, + * disclosure or distribution of this material and related documentation + * without an express license agreement from NVIDIA CORPORATION or + * its affiliates is strictly prohibited. + */ + +#include "tensorrt_llm/executor/transferAgent.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/interfaces.h" +#include +#include + +using namespace tensorrt_llm::executor::kv_cache; + +class LocalAgentRegistrar final : public tensorrt_llm::executor::kv_cache::AgentRegistrar +{ +public: + [[nodiscard]] AgentDesc const* getAgentDesc(char const* agentName) const + { + auto it = mAgentDescs.find(agentName); + TLLM_CHECK(it != mAgentDescs.end()); + return &it->second; + } + + void addAgentDesc(char const* agentName, AgentDesc desc) + { + mAgentDescs.insert(std::make_pair(agentName, std::move(desc))); + } + + void removeAgentDesc(char const* agentName) + { + mAgentDescs.erase(agentName); + } + +private: + std::unordered_map mAgentDescs; +}; + +class RegisteredHostMemory +{ +public: + RegisteredHostMemory(MemoryDescs mems, BaseTransferAgent* agent) + : mDescs{std::move(mems)} + , mAgentPtr{agent} + { + TLLM_CHECK(mAgentPtr); + mAgentPtr->registerMemory(mDescs); + } + + ~RegisteredHostMemory() + { + TLLM_CHECK(mAgentPtr); + mAgentPtr->deregisterMemory(mDescs); + } + + [[nodiscard]] MemoryDescs const& getDescs() const noexcept + { + return mDescs; + } + +private: + MemoryDescs mDescs; + BaseTransferAgent* mAgentPtr{}; +}; + +class TransferAgentTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) +{ +public: + void SetUp() override {} + + void TearDown() override {} + + [[nodiscard]] std::unique_ptr makeTransferAgent( + BaseAgentConfig const& config, AgentRegistrar* registrar) + { + return tensorrt_llm::executor::kv_cache::createNixlTransferAgent(&config, registrar); + } +}; + +TEST_F(TransferAgentTest, Basic) +{ + LocalAgentRegistrar registrar; + + char const *agent0{"agent0"}, *agent1{"agent1"}; + BaseAgentConfig config0{agent0, true}, config1{agent1, true}; + auto nixlAgent0 = makeTransferAgent(config0, ®istrar); + auto nixlAgent1 = makeTransferAgent(config1, ®istrar); + + TLLM_CHECK(nixlAgent0); + TLLM_CHECK(nixlAgent1); + + std::vector memory0(100, 10); + std::vector memory1(100, 1); + + RegisteredHostMemory regMem0(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory0}}}, nixlAgent0.get()); + RegisteredHostMemory regMem1(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, nixlAgent1.get()); + + nixlAgent0->loadRemoteAgent(agent1); + + TransferRequest writeReq{TransferOp::kWRITE, regMem0.getDescs(), regMem1.getDescs(), agent1}; + auto status = nixlAgent0->submitTransferRequests(writeReq); + status->wait(); + + TLLM_CHECK(memory0 == memory1); + + nixlAgent0->invalidateRemoteAgent(agent1); +} diff --git a/docker/common/install_nixl.sh b/docker/common/install_nixl.sh index cd31e07ee3ba..91110c3af801 100644 --- a/docker/common/install_nixl.sh +++ b/docker/common/install_nixl.sh @@ -5,6 +5,7 @@ GITHUB_URL="https://github.com" UCX_VERSION="v1.18.1" UCX_INSTALL_PATH="/usr/local/ucx/" +CUDA_PATH="/usr/local/cuda" NIXL_VERSION="0.2.0" @@ -23,7 +24,18 @@ if [ ! -d ${UCX_INSTALL_PATH} ]; then git clone --depth 1 -b ${UCX_VERSION} ${UCX_REPO} cd ucx ./autogen.sh - ./contrib/configure-release --prefix=${UCX_INSTALL_PATH} + ./contrib/configure-release \ + --prefix=${UCX_INSTALL_PATH} \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --with-cuda=${CUDA_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt make install -j$(nproc) cd .. rm -rf ucx # Remove UCX source to save space @@ -36,12 +48,17 @@ if [ "$(uname -m)" != "amd64" ] && [ "$(uname -m)" != "x86_64" ]; then EXTRA_NIXL_ARGS="-Ddisable_gds_backend=true" fi +if [ $ARCH_NAME != "x86_64-linux-gnu" ]; then + echo "The NIXL backend is temporarily unavailable on the aarch64 platform. Exiting script." + exit 0 +fi + pip3 install --no-cache-dir meson ninja pybind11 git clone --depth 1 -b ${NIXL_VERSION} ${NIXL_REPO} cd nixl -meson setup builddir -Ducx_path=${UCX_INSTALL_PATH} ${EXTRA_NIXL_ARGS} +meson setup builddir -Ducx_path=${UCX_INSTALL_PATH} -Dstatic_plugins=UCX ${EXTRA_NIXL_ARGS} cd builddir && ninja install cd ../.. -rm -rf nixl # Remove NIXL source tree to save space +rm -rf nixl* # Remove NIXL source tree to save space echo "export LD_LIBRARY_PATH=/opt/nvidia/nvda_nixl/lib/${ARCH_NAME}:/opt/nvidia/nvda_nixl/lib64:\$LD_LIBRARY_PATH" >> "${ENV}" diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index c7c3c6d0a068..bac98b5d9904 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -50,7 +50,7 @@ def BUILD_CONFIGS = [ // Vanilla TARNAME is used for packaging in runLLMPackage // cmake-vars cannot be empty, so passing (default) multi-device configuration. (CONFIG_LINUX_X86_64_VANILLA) : [ - (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=1 --extra-cmake-vars WARNING_IS_ERROR=ON --micro_benchmarks", + (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=1 --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl --micro_benchmarks", (TARNAME) : "TensorRT-LLM.tar.gz", (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;120-real", ], diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 86325b71313d..430cf1093428 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -21,10 +21,10 @@ UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifac // Container configuration // available tags can be found in: https://urm.nvidia.com/artifactory/sw-tensorrt-docker/tensorrt-llm/ // [base_image_name]-[arch]-[os](-[python_version])-[trt_version]-[torch_install_type]-[stage]-[date]-[mr_id] -LLM_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505121727-4049" -LLM_SBSA_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-aarch64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505121727-4049" -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.10.0.31-skip-tritondevel-202505131825-4114" -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.10.0.31-skip-tritondevel-202505131825-4114" +LLM_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505160532-3934" +LLM_SBSA_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-aarch64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505160532-3934" +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.10.0.31-skip-tritondevel-202505160532-3934" +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.10.0.31-skip-tritondevel-202505160532-3934" // TODO: Move common variables to an unified location BUILD_CORES_REQUEST = "8" diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index a65fa5319c25..767069cb71d3 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -483,6 +483,14 @@ def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): build_dir / "tensorrt_llm/executor/cache_transmission/ucx_utils/libtensorrt_llm_ucx_wrapper.so", lib_dir / "libtensorrt_llm_ucx_wrapper.so") + if os.path.exists( + build_dir / + "tensorrt_llm/executor/cache_transmission/nixl_utils/libtensorrt_llm_nixl_wrapper.so" + ): + install_file( + build_dir / + "tensorrt_llm/executor/cache_transmission/nixl_utils/libtensorrt_llm_nixl_wrapper.so", + lib_dir / "libtensorrt_llm_nixl_wrapper.so") install_file( build_dir / "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/libdecoder_attention_0.so", diff --git a/scripts/rename_docker_images.py b/scripts/rename_docker_images.py index b844afdf5b1c..405deb69bf7a 100755 --- a/scripts/rename_docker_images.py +++ b/scripts/rename_docker_images.py @@ -55,7 +55,7 @@ def parse_arguments() -> _ap.Namespace: def get_current_timestamp() -> str: """Get the current timestamp in YYYYMMDDhhmm format.""" - return _dt.datetime.now(_dt.UTC).strftime("%Y%m%d%H%M") + return _dt.datetime.now(_dt.timezone.utc).strftime("%Y%m%d%H%M") def run_shell_command(command: str, dry_run: bool) -> None: @@ -209,7 +209,7 @@ def rename_images(*, run_shell_command(f"docker push {dst_image}", dry_run) find_and_replace_in_files(base_dir / "jenkins", ".groovy", dst_image_old, dst_image, dry_run) - find_and_replace_in_files(base_dir / ".devcontainer", ".yaml", + find_and_replace_in_files(base_dir / ".devcontainer", ".yml", dst_image_old, dst_image, dry_run) diff --git a/setup.py b/setup.py index af1c8e244b3b..1e28393cedcf 100644 --- a/setup.py +++ b/setup.py @@ -98,6 +98,7 @@ def has_ext_modules(self): 'bin/executorWorker', 'libs/libtensorrt_llm.so', 'libs/libth_common.so', 'libs/libnvinfer_plugin_tensorrt_llm.so', 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', + 'libs/libtensorrt_llm_nixl_wrapper.so', 'libs/libdecoder_attention_1.so', 'bindings.*.so', "include/**/*" ] diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index ee9860a6ce17..8b0c950dd8c3 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -76,6 +76,20 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): env=new_env, timeout=600) + new_env = copy.copy(cpp_env) + # Transfer agent tests + transfer_agent_test_2_proc = [ + "mpirun", + "-n", + "2", + "--allow-run-as-root", + "executor/transferAgentTest", + ] + _cpp.run_command(transfer_agent_test_2_proc, + cwd=tests_dir, + env=new_env, + timeout=600) + def run_llama_executor_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): tests_dir = build_dir / "tests"