[None][doc] Facilitates the integration of the transfer agent (#7867)

Signed-off-by: Shixiaowei02 <[email protected]>

Author: Shixiaowei02

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -40,6 +40,8 @@ enum class MemoryType : uint8_t
     kFILE
 };
 
+// `MemoryDesc` is used to describe a memory region, which can then be designated
+// as the source or destination of read/write operations.
 class MemoryDesc
 {
 public:
@@ -192,6 +194,8 @@ using RegisterDescs = MemoryDescs;
 using SyncMessage = std::string;
 using ConnectionInfoType = std::string;
 
+// `AgentDesc` represents the unique identifier for reading and writing to the agent.
+// By accessing this identifier, the backend can establish the correct connection.
 class AgentDesc final
 {
 public:
@@ -209,15 +213,24 @@ class AgentDesc final
     std::string mBackendAgentDesc;
 };
 
+// `TransferOp` is an enumeration that represents the types of transfer operations.
+// Currently, it supports two operations: `read` and `write`.
 enum class TransferOp : uint8_t
 {
     kREAD,
     kWRITE,
 };
 
+// `TransferRequest` is used to represent the transfer requests supported by the underlying agent.
 class TransferRequest
 {
 public:
+    /// @brief The constructor of `TransferRequest`.
+    /// @param op Source data arrangement.
+    /// @param srcDescs Description of the source memory region.
+    /// @param dstDescs Description of the destination memory region.
+    /// @param remoteName Name of the remote counterpart.
+    /// @param syncMessage Synchronization information for the end of the transfer.
     TransferRequest(TransferOp op, TransferDescs srcDescs, TransferDescs dstDescs, std::string const& remoteName,
         std::optional<SyncMessage> syncMessage = std::nullopt)
         : mOp{op}
@@ -261,6 +274,7 @@ class TransferRequest
     std::optional<SyncMessage> mSyncMessage;
 };
 
+// Data structure for checking the status of active transfer operations.
 class TransferStatus
 {
 public:
@@ -281,22 +295,52 @@ class BaseTransferAgent
 public:
     virtual ~BaseTransferAgent() = default;
 
+    /// @brief Register a memory region.
+    /// @param descs Describe the memory regions to be registered.
     virtual void registerMemory(RegisterDescs const& descs) = 0;
 
+    /// @brief Unregister a memory region.
+    /// @param descs Describe the memory regions to be unregistered.
     virtual void deregisterMemory(RegisterDescs const& descs) = 0;
 
+    /// @brief Initialize and establish a connection with a remote agent.
+    /// @param name Specify the name of the remote agent.
+    /// @param agentDesc Provide the necessary communication details for connecting to the remote agent.
     virtual void loadRemoteAgent(std::string const& name, AgentDesc const& agentDesc) = 0;
-    virtual AgentDesc getLocalAgentDesc() = 0;
 
+    /// @brief Initialize and establish a connection with a remote agent.
+    /// @param name Specify the name of the remote agent.
+    /// @param connectionInfo Provide the necessary communication details for connecting to the remote agent.
+    virtual void loadRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) = 0;
+
+    /// @brief Invalidate a connection with a remote agent.
+    /// @param name Specify the name of the remote agent.
     virtual void invalidateRemoteAgent(std::string const& name) = 0;
 
+    /// @brief Fetch the descriptor of the local agent.
+    /// @return The descriptor of the local agent.
+    virtual AgentDesc getLocalAgentDesc() = 0;
+
+    /// @brief Fetch the descriptor of the local agent.
+    /// @return The descriptor of the local agent.
+    virtual ConnectionInfoType getLocalConnectionInfo() = 0;
+
+    /// @brief Initiate the transfer by submitting the request.
+    /// @param request Specify the transmission request.
+    /// @return The status of the requests.
     [[nodiscard]] virtual std::unique_ptr<TransferStatus> submitTransferRequests(TransferRequest const& request) = 0;
+
+    /// @brief Generate a notification, not bound to a transfer, e.g., for control.
+    /// @param name Specify the name of the remote agent to which the information should be sent.
+    /// @param syncMessage The data or message intended for synchronization.
     virtual void notifySyncMessage(std::string const& name, SyncMessage const& syncMessage) = 0;
 
+    /// @brief Retrieve notification messages sent by other agents.
+    /// @return A mapping from remote agent names to their respective notification messages.
     virtual std::unordered_map<std::string, std::vector<SyncMessage>> getNotifiedSyncMessages() = 0;
 
-    virtual ConnectionInfoType getConnectionInfo() = 0;
-    virtual void connectRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) = 0;
+    /// @brief Check if metadata is available for a remote agent.
+    /// @return Whether the metadata is available for a remote agent.
     virtual bool checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) = 0;
 };
 

cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

@@ -145,16 +145,18 @@ using runtime::SizeType32;
 using AgentConnectionManager = tensorrt_llm::executor::kv_cache::AgentConnectionManager;
 using DataContext = tensorrt_llm::executor::kv_cache::DataContext;
 
-static int32_t tagFromRequestId(LlmRequest::RequestIdType requestId)
+namespace
+{
+
+int32_t tagFromRequestId(LlmRequest::RequestIdType requestId)
 {
     constexpr int32_t kDATA_TAG{43};
     return ((requestId & 0xFFF) << 8) | (kDATA_TAG & 0xFF);
 }
 
-namespace fs = std::filesystem;
-
-static fs::path getTransferOutputPath(char const* tag)
+std::filesystem::path getTransferOutputPath(char const* tag)
 {
+    namespace fs = std::filesystem;
     auto outputPath = common::getEnvKVCacheTransferOutputPath();
     if (!outputPath.empty())
     {
@@ -166,13 +168,15 @@ static fs::path getTransferOutputPath(char const* tag)
     return {};
 }
 
+} // namespace
+
 struct ReceiveCacheResource
 {
     runtime::BufferManager mBufferManager;
     runtime::CudaEvent mCudaEvent;
 
-    ReceiveCacheResource(runtime::BufferManager&& bufferManager, runtime::CudaEvent&& cudaEvent)
-        : mBufferManager(bufferManager)
+    ReceiveCacheResource(runtime::BufferManager&& bufferManager, runtime::CudaEvent cudaEvent)
+        : mBufferManager(std::move(bufferManager))
         , mCudaEvent(std::move(cudaEvent))
     {
     }
@@ -343,8 +347,7 @@ class CacheSender::Impl
         TLLM_CHECK_WITH_INFO(mFormatter->inquireSupport(
                                  mSelfState.getCacheState().value(), info.getTransState().getCacheState().value()),
             "Disagg server does not currently support these cacheState, please check the cacheState of the context and "
-            "gen "
-            "executors");
+            "gen executors");
         auto peerRelativeRanks = executor::kv_cache::targetIRanks(info.getTransState().getCacheState().value(),
             mSelfState.getCacheState().value(), mSelfState.getCommState().value().getSelfIdx())
                                      .mIRanks;
@@ -1024,7 +1027,6 @@ class CacheReceiver::Impl
 
     void request(AsyncResource& resource)
     {
-
         tensorrt_llm::common::setThreadName("dataTransRequest");
         TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId));
 

cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp

@@ -144,7 +144,7 @@ void AgentConnection::sendRequestAndBufferInfo(
     TLLM_CHECK(deviceId == mAgentConnectionManager->getDeviceId());
     MemoryDesc bufferDesc(
         reinterpret_cast<uintptr_t>(preAllocateBuffer->data()), preAllocateBuffer->getSize(), deviceId);
-    std::string address = mAgentConnectionManager->getAgent()->getConnectionInfo();
+    std::string address = mAgentConnectionManager->getAgent()->getLocalConnectionInfo();
     std::optional<std::string> metadataOpt = std::nullopt;
     if (mNeedSendMetadata)
     {
@@ -225,7 +225,7 @@ AgentConnectionManager::AgentConnectionManager(
     mRegMemDescs = MemoryDescs{MemoryType::kVRAM, MemDescs};
     m_Agent->registerMemory(mRegMemDescs);
 
-    AgentState localAgentState{mAgentName, m_Agent->getConnectionInfo()};
+    AgentState localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()};
     std::vector<AgentState> agentStates(mpi::MpiComm::session().getSize());
     if (mpi::MpiComm::session().getSize() > 1)
     {
@@ -411,10 +411,10 @@ AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentN
         }
         else
         {
-            TLLM_CHECK_WITH_INFO(!isSender, "Sender shouldn't call connectRemoteAgent");
-            TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "mAgentName: %s connect to %s with connectRemoteAgent",
+            TLLM_CHECK_WITH_INFO(!isSender, "Sender shouldn't call loadRemoteAgent");
+            TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "mAgentName: %s connect to %s with loadRemoteAgent",
                 mAgentName.c_str(), remoteAgentName.c_str());
-            m_Agent->connectRemoteAgent(remoteAgentName, connectionInfo);
+            m_Agent->loadRemoteAgent(remoteAgentName, connectionInfo);
         }
     }
     else

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp

@@ -469,19 +469,19 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c
     return notifs;
 }
 
-ConnectionInfoType NixlTransferAgent::getConnectionInfo()
+ConnectionInfoType NixlTransferAgent::getLocalConnectionInfo()
 {
     return mAddress;
 }
 
-void NixlTransferAgent::connectRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo)
+void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo)
 {
     std::string ip = connectionInfo.substr(0, connectionInfo.find(":"));
     std::string port = connectionInfo.substr(connectionInfo.find(":") + 1);
     TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
-        "NixlTransferAgent::connectRemoteAgent connectRemoteAgent to %s remoteagent name: %s", connectionInfo.c_str(),
+        "NixlTransferAgent::loadRemoteAgent loadRemoteAgent to %s remoteagent name: %s", connectionInfo.c_str(),
         name.c_str());
-    TLLM_CHECK_WITH_INFO(!ip.empty() && !port.empty(), "connectRemoteAgent get empty ip or port, connectionInfo: %s",
+    TLLM_CHECK_WITH_INFO(!ip.empty() && !port.empty(), "loadRemoteAgent get empty ip or port, connectionInfo: %s",
         connectionInfo.c_str());
     nixl_opt_args_t md_extra_params;
     md_extra_params.ipAddr = ip;
@@ -506,7 +506,7 @@ void NixlTransferAgent::connectRemoteAgent(std::string const& name, ConnectionIn
         }
     }
     TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
-        "NixlTransferAgent::connectRemoteAgent connectRemoteAgent to %s remoteagent name: %s success status: %s",
+        "NixlTransferAgent::loadRemoteAgent loadRemoteAgent to %s remoteagent name: %s success status: %s",
         connectionInfo.c_str(), name.c_str(), nixlEnumStrings::statusStr(status).c_str());
 }
 

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h

@@ -84,9 +84,9 @@ class NixlTransferAgent final : public BaseTransferAgent
 
     [[nodiscard]] std::unordered_map<std::string, std::vector<SyncMessage>> getNotifiedSyncMessages() override;
 
-    ConnectionInfoType getConnectionInfo() override;
+    ConnectionInfoType getLocalConnectionInfo() override;
 
-    void connectRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) override;
+    void loadRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) override;
 
     bool checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) override;
 

cpp/tests/unit_tests/executor/transferAgentTest.cpp

@@ -80,8 +80,8 @@ TEST_F(TransferAgentTest, Basic)
     RegisteredHostMemory regMem1(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, nixlAgent1.get());
 
     // nixlAgent0->loadRemoteAgent(agent1);
-    auto connectionInfo = nixlAgent1->getConnectionInfo();
-    nixlAgent0->connectRemoteAgent(agent1, connectionInfo);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
     bool checked = false;
     do
     {
@@ -116,8 +116,8 @@ TEST_F(TransferAgentTest, Basic2)
     RegisteredHostMemory regMem1(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, nixlAgent1.get());
 
     // nixlAgent0->loadRemoteAgent(agent1);
-    auto connectionInfo = nixlAgent1->getConnectionInfo();
-    nixlAgent0->connectRemoteAgent(agent1, connectionInfo);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
     bool checked = false;
     do
     {
@@ -159,8 +159,8 @@ TEST_F(TransferAgentTest, DeviceMemory)
         MemoryDescs{MemoryType::kVRAM, {MemoryDesc{dev_ptr1, size, deviceId}}}, nixlAgent1.get());
 
     // nixlAgent0->loadRemoteAgent(agent1);
-    auto connectionInfo = nixlAgent1->getConnectionInfo();
-    nixlAgent0->connectRemoteAgent(agent1, connectionInfo);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
     bool checked = false;
     do
     {
@@ -201,8 +201,8 @@ TEST_F(TransferAgentTest, Connect)
     nixlAgent2->registerMemory(memDescs0);
 
     // nixlAgent0->loadRemoteAgent(agent1);
-    auto connectionInfo = nixlAgent1->getConnectionInfo();
-    nixlAgent0->connectRemoteAgent(agent1, connectionInfo);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
     bool checked = false;
     do
     {
@@ -213,7 +213,7 @@ TEST_F(TransferAgentTest, Connect)
     status->wait();
 
     TLLM_CHECK(memory0 == memory1);
-    nixlAgent2->connectRemoteAgent(agent1, connectionInfo);
+    nixlAgent2->loadRemoteAgent(agent1, connectionInfo);
     checked = false;
     do
     {
@@ -251,8 +251,8 @@ TEST_F(TransferAgentTest, SyncMessage)
     RegisteredHostMemory regMem3(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, nixlAgent1.get());
 
     // nixlAgent0->loadRemoteAgent(agent1);
-    auto connectionInfo = nixlAgent1->getConnectionInfo();
-    nixlAgent0->connectRemoteAgent(agent1, connectionInfo);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
     bool checked = false;
     do
     {
@@ -287,8 +287,8 @@ TEST_F(TransferAgentTest, SyncMessage)
     TLLM_CHECK(notif2[agent0][0] == syncMessage2);
 
     // nixlAgent1->loadRemoteAgent(agent0);
-    auto connectionInfo2 = nixlAgent0->getConnectionInfo();
-    nixlAgent1->connectRemoteAgent(agent0, connectionInfo2);
+    auto connectionInfo2 = nixlAgent0->getLocalConnectionInfo();
+    nixlAgent1->loadRemoteAgent(agent0, connectionInfo2);
     std::string syncMessage3 = "three_agent_sync_message";
     nixlAgent1->notifySyncMessage(agent0, syncMessage3);
     auto notif3 = nixlAgent0->getNotifiedSyncMessages();

docs/source/developer-guide/kv-transfer.md

@@ -0,0 +1,70 @@
+# Introduction to KV Cache Transmission
+
+This article provides a general overview of the components used for device-to-device transmission of KV cache, which is relied upon by dist-serving. It is intended as a reference for users who wish to understand the internal implementation or develop extended functionalities.
+
+## Table of Contents
+
+- [Workflow](#workflow)
+- [Key Components](#key-components)
+  - [Transceiver](#transceiver)
+  - [Sender and Receiver](#sender-and-receiver)
+  - [Formatter](#formatter)
+  - [Connection](#connection)
+  - [Transfer Agent](#transfer-agent)
+- [Customization](#customization)
+  - [Encapsulation and Overloading of Low-Level Communication Libraries](#encapsulation-and-overloading-of-low-level-communication-libraries)
+  - [Modifications to Upper-Level Runtime Logic](#modifications-to-upper-level-runtime-logic)
+- [Evolution Outlook](#evolution-outlook)
+
+## Workflow
+
+<img src="https://github.com/NVIDIA/TensorRT-LLM/blob/rel/docs/source/media/kv_transfer.png?raw=true" alt="KV Cache Transfer Overview" width="500" height="auto">
+
+1. Context phase completes computation, KV cache stays in device memory awaiting transmission.
+2. Context returns its communicator handle to the user, who selects the generation executor for continued communication.
+3. If no prior connection exists, it's established now. Generation phase shares its cache layout with context.
+4. Generation phase requests KV cache for specific tokens.
+5. Context sends KV cache to generation phase.
+6. Generation phase resumes computation, context releases KV cache.
+
+## Key Components
+
+### Transceiver
+
+Responsible for coordinating the sending and receiving of cache among different ranks within the same executor.
+
+### Sender and Receiver
+
+Responsible for transmitting control plane messages. That is, during per-request transmission, the receiver bound to the generation informs the sender of the specific information it requires. The sender then sends the corresponding KV cache based on these messages.
+
+### Formatter
+
+Performs KV cache data transmission and correctly handles the mapping between caches across different TP/PP configurations.
+
+### Connection
+
+Bidirectional byte-stream protocol facility. Apart from essential operations such as connection establishment, it mainly provides send and receive functionalities. UCX accesses the system through this facility. The `AgentConnection` data structure adapts the upper-layer bidirectional send/receive semantics into a unidirectional read/write operation model.
+
+### Transfer Agent
+
+Unidirectional byte-stream read/write protocol facility. Apart from essential operations such as connection establishment, it primarily provides read and write functionalities. NIXL accesses the system through this facility.
+
+## Customization
+
+At the current stage, the customization work mainly involves inheriting the low-level data plane interfaces to enable the invocation of third-party communication libraries, as well as defining the data structures required for establishing connections in the data plane.
+
+### Encapsulation and Overloading of Low-Level Communication Libraries
+
+Each layer of interface described in the previous section supports overloading. Here, based on whether the underlying library uses a unidirectional or bidirectional protocol, we describe the customization methods respectively.
+
+If the underlying library you are integrating uses a unidirectional communication model, with read/write as its primary interfaces, you should inherit the `executor::kv_cache::BaseTransferAgent` data structure. This structure mainly provides interfaces for memory registration, remote agent loading, and transfer request submission.
+
+If the underlying library you are integrating uses a bidirectional communication model, you should inherit the `executor::kv_cache::Connection` data structure. This structure mainly provides send and receive interfaces.
+
+### Modifications to Upper-Level Runtime Logic
+
+This corresponds to the communication info section shown in the figure above. Since different underlying communication connections may require completely different setup methods—for example, some use IP and port, others require a world rank, and some communication libraries establish connections using binary-transparent metadata—we provide sufficient flexibility to allow users to customize this part as needed.
+
+## Evolution Outlook
+
+Currently, the architecture of KV transfer is being optimized. First, we plan to move the control plane logic up to Python to enable better integration with the Python runtime. In addition, we are reevaluating the current design choice of initiating communication only after the context computation is completed, which was originally made for flexibility. Lastly, since some control logic is still being transmitted through the data plane, we aim to clarify the relationship between the control and data planes, and to simplify and streamline the code logic of the data plane. Due to the modular architecture, these iterative enhancements are only loosely coupled with the `TransferAgent`. We aim to minimize the impact of future upgrades on third-party integrations.

docs/source/index.rst

@@ -85,6 +85,7 @@ Welcome to TensorRT LLM's Documentation!
    developer-guide/ci-overview.md
    developer-guide/dev-containers.md
    developer-guide/api-change.md
+   developer-guide/kv-transfer.md
 
 
 .. toctree::