feat: NIXL interface integration

.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
 

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()

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<bool>(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<bool>(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)
 

cpp/include/tensorrt_llm/executor/cacheCommunicator.h

@@ -16,21 +16,14 @@
 
 #pragma once
 
-#include <future>
-#include <variant>
-
-#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 <cstddef>
+#include <vector>
 
 namespace tensorrt_llm::executor::kv_cache
 {
 
+class CommState;
+
 struct DataContext
 {
 public:

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 <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace tensorrt_llm::executor::kv_cache
+{
+
+enum class MemoryType : uint8_t
+{
+    kDRAM,
+    kVRAM,
+    kBLK,
+    kOBJ,
+    kFILE
+};
+
+class MemoryDesc
+{
+public:
+    MemoryDesc(std::vector<char> const& vec, uint32_t deviceId = 0)
+        : mAddr{reinterpret_cast<uintptr_t>(vec.data())}
+        , mLen{vec.size()}
+        , mDeviceId{deviceId}
+    {
+    }
+
+    MemoryDesc(void* addr, size_t len, uint32_t deviceId)
+        : mAddr{reinterpret_cast<uintptr_t>(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<MemoryDesc> descs)
+        : mType{type}
+        , mDescs{std::move(descs)}
+    {
+    }
+
+    [[nodiscard]] MemoryType getType() const noexcept
+    {
+        return mType;
+    }
+
+    [[nodiscard]] std::vector<MemoryDesc> const& getDescs() const noexcept
+    {
+        return mDescs;
+    }
+
+private:
+    MemoryType mType;
+    std::vector<MemoryDesc> mDescs;
+};
+
+using TransferDescs = MemoryDescs;
+using RegisterDescs = MemoryDescs;
+
+class AgentDesc final
+{
+public:
+    AgentDesc(std::vector<char> backendAgentDesc)
+        : mBackendAgentDesc{backendAgentDesc}
+    {
+    }
+
+    [[nodiscard]] std::vector<char> const& getBackendAgentDesc() const noexcept
+    {
+        return mBackendAgentDesc;
+    }
+
+private:
+    std::vector<char> 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<TransferStatus> 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 <typename FunctionT>
+    [[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<FunctionT>(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<std::string, void*> mHandlers;
+};
+
+template <typename... Args>
+[[nodiscard]] std::unique_ptr<BaseTransferAgent> makeTransferAgent(char const* const& backend, Args&&... args)
+{
+    if (backend == std::string{"nixl"})
+    {
+        auto& loader = DynLibLoader::getInstance();
+        using CreateNixlFuncType = std::unique_ptr<BaseTransferAgent> (*)(BaseAgentConfig const*, AgentRegistrar*);
+        auto* func = loader.getFunctionPointer<CreateNixlFuncType>(
+            "libtensorrt_llm_nixl_wrapper.so", "createNixlTransferAgent");
+        return func(std::forward<Args>(args)...);
+    }
+    TLLM_THROW("Unknown backend name.");
+}
+
+} // namespace tensorrt_llm::executor::kv_cache