[None][feat] support NIXL cache transceiver with Ray

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -381,6 +381,8 @@ struct BaseAgentConfig
     bool useListenThread;
     bool enableTelemetry;
     std::unordered_map<std::string, std::string> backendParams;
+    std::optional<int> rank;
+    std::optional<int> worldSize;
 };
 
 class BaseTransferAgent

cpp/tensorrt_llm/common/envUtils.cpp

@@ -511,12 +511,6 @@ bool getEnvKVCachePoolUseFabricMemory()
     return useFabricMemory;
 }
 
-uint16_t getEnvNixlPort()
-{
-    static uint16_t const nixlPort = getUInt64Env("TRTLLM_NIXL_PORT").value_or(0);
-    return nixlPort;
-}
-
 bool getEnvNixlDisableCoalesce()
 {
     static bool const disableCoalesce = getBoolEnv("TRTLLM_NIXL_DISABLE_COALESCE");

cpp/tensorrt_llm/common/envUtils.h

@@ -149,8 +149,6 @@ size_t getEnvMemSizeForKVCacheTransferBuffer();
 
 bool getEnvKVCachePoolUseFabricMemory();
 
-uint16_t getEnvNixlPort();
-
 // Whether to disable coalescing of contiguous NIXL transfer descriptors (coalescing is on by default).
 bool getEnvNixlDisableCoalesce();
 

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

@@ -18,12 +18,17 @@
 #include "connection.h"
 #include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h"
+#include "tensorrt_llm/runtime/utils/pgUtils.h"
 #include <limits>
+#include <numeric>
 #include <random>
 #include <string>
 #include <unistd.h>
 #include <utility>
 
+using tensorrt_llm::pg_utils::get_world_pg;
+using tensorrt_llm::pg_utils::PgHelper;
+
 namespace tensorrt_llm::executor::kv_cache
 {
 
@@ -390,9 +395,36 @@ AgentConnectionManager::AgentConnectionManager(
     TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId));
     TLLM_CHECK(mDeviceId != -1);
 
+    c10::intrusive_ptr<c10d::ProcessGroup> worldPg;
+    if (useMPI())
+    {
+        mRank = mpi::MpiComm::session().getRank();
+        mWorldSize = mpi::MpiComm::session().getSize();
+    }
+    else
+    {
+        worldPg = get_world_pg();
+        if (worldPg)
+        {
+            mRank = worldPg->getRank();
+            mWorldSize = worldPg->getSize();
+            TLLM_LOG_DEBUG(
+                mRank, "Cache transceiver using Torch process group - rank: %d, world size: %d", mRank, mWorldSize);
+        }
+        else
+        {
+            TLLM_LOG_WARNING(
+                "Torch process group is not initialized while MPI is disabled; cache transceiver defaults to one "
+                "process. For multi-rank execution, initialize the Torch process group before constructing the cache "
+                "transceiver, or unset TLLM_DISABLE_MPI to use MPI");
+        }
+    }
+
     mAgentName = genUniqueAgentName();
     // Create Agent
     BaseAgentConfig config{mAgentName, true, false, true};
+    config.rank = mRank;
+    config.worldSize = mWorldSize;
     m_Agent = makeTransferAgent(backendType, &config);
     TLLM_CHECK(!mCacheTransBufferManagers.empty());
     mBufferKinds.reserve(mCacheTransBufferManagers.size());
@@ -418,49 +450,71 @@ AgentConnectionManager::AgentConnectionManager(
     m_Agent->registerMemory(mRegMemDescs);
 
     AgentState localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()};
-    std::vector<AgentState> agentStates(mpi::MpiComm::session().getSize());
-    if (mpi::MpiComm::session().getSize() > 1)
+    std::vector<AgentState> agentStates(mWorldSize);
+    if (mWorldSize > 1)
     {
-
-        mpi::MpiComm::session().barrier();
         namespace su = executor::serialize_utils;
 
         std::ostringstream oStream;
         su::serialize(localAgentState, oStream);
         auto str = oStream.str();
         std::vector<char> buffer(str.begin(), str.end());
-        std::vector<SizeType32> sizeofBuffer(mpi::MpiComm::session().getSize());
+        std::vector<SizeType32> sizeofBuffer(mWorldSize);
         SizeType32 bufferSize = buffer.size();
-        mpi::MpiComm::session().allgather(&bufferSize, sizeofBuffer.data(), 1, mpi::MpiType::kINT32);
-        SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0);
-        std::vector<char> recvBuffer(recvBufferSize);
-        std::vector<int> displs(mpi::MpiComm::session().getSize());
-        for (int r = 0; r < mpi::MpiComm::session().getSize(); r++)
+
+        if (useMPI())
         {
-            displs[r] = (r == 0) ? 0 : (displs[r - 1] + sizeofBuffer[r - 1]);
-        }
-        mpi::MpiComm::session().allgatherv(buffer.data(), bufferSize, mpi::MpiType::kCHAR, recvBuffer.data(),
-            sizeofBuffer, displs, mpi::MpiType::kCHAR);
+            mpi::MpiComm::session().barrier();
+            mpi::MpiComm::session().allgather(&bufferSize, sizeofBuffer.data(), 1, mpi::MpiType::kINT32);
+            SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0);
+            std::vector<char> recvBuffer(recvBufferSize);
+            std::vector<int> displs(mWorldSize);
+            for (int r = 0; r < mWorldSize; r++)
+            {
+                displs[r] = (r == 0) ? 0 : (displs[r - 1] + sizeofBuffer[r - 1]);
+            }
+            mpi::MpiComm::session().allgatherv(buffer.data(), bufferSize, mpi::MpiType::kCHAR, recvBuffer.data(),
+                sizeofBuffer, displs, mpi::MpiType::kCHAR);
 
-        // deserialize
-        for (int i = 0; i < mpi::MpiComm::session().getSize(); i++)
+            for (int r = 0; r < mWorldSize; r++)
+            {
+                std::vector<char> serBuffer(
+                    recvBuffer.begin() + displs[r], recvBuffer.begin() + (displs[r] + sizeofBuffer[r]));
+                su::VectorWrapBuf<char> strbuf(serBuffer);
+                std::istream is(&strbuf);
+                agentStates[r] = su::deserialize<executor::kv_cache::AgentState>(is);
+                TLLM_LOG_DEBUG(mRank, " recv agentStates[%d]: %s", r, agentStates[r].toString().c_str());
+            }
+        }
+        else
         {
-            std::vector<char> serBuffer(
-                recvBuffer.begin() + displs[i], recvBuffer.begin() + (displs[i] + sizeofBuffer[i]));
-            su::VectorWrapBuf<char> strbuf(serBuffer);
-            std::istream is(&strbuf);
-            agentStates[i] = su::deserialize<executor::kv_cache::AgentState>(is);
-            TLLM_LOG_DEBUG(
-                mpi::MpiComm::world().getRank(), " recv  agentStates[%d]: %s", i, agentStates[i].toString().c_str());
+            PgHelper pgHelper{worldPg};
+            PGCHECK_THROW(worldPg->barrier());
+            PGCHECK_THROW(pgHelper.allgather(&bufferSize, std::ref(sizeofBuffer), {}));
+
+            SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0);
+            std::vector<char> recvBuffer(recvBufferSize);
+            PGCHECK_THROW(pgHelper.allgatherv(std::ref(buffer), std::ref(recvBuffer), std::cref(sizeofBuffer), {}));
+
+            char* begin = recvBuffer.data();
+            for (int r = 0; r < mWorldSize; r++)
+            {
+                std::vector<char> serBuffer(begin, begin + sizeofBuffer[r]);
+                begin += sizeofBuffer[r];
+                su::VectorWrapBuf<char> strbuf(serBuffer);
+                std::istream is(&strbuf);
+                agentStates[r] = su::deserialize<executor::kv_cache::AgentState>(is);
+                TLLM_LOG_DEBUG(mRank, " recv agentStates[%d]: %s", r, agentStates[r].toString().c_str());
+            }
         }
     }
     else
     {
         agentStates[0] = localAgentState;
     }
-    mCommState = CommState(agentStates, mpi::MpiComm::session().getRank());
-    TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
-        " ***** AgentConnectionManager::AgentConnectionManager    mCommState: %s", mCommState.toString().c_str());
+    mCommState = CommState(agentStates, mRank);
+    TLLM_LOG_DEBUG(
+        mRank, " ***** AgentConnectionManager::AgentConnectionManager mCommState: %s", mCommState.toString().c_str());
 }
 
 AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo(
@@ -698,8 +752,7 @@ AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentN
     std::optional<std::string> metadata, bool isSender)
 {
 
-    TLLM_LOG_DEBUG(
-        mpi::MpiComm::world().getRank(), "mAgentName: %s connect to %s", mAgentName.c_str(), remoteAgentName.c_str());
+    TLLM_LOG_DEBUG(mRank, "mAgentName: %s connect to %s", mAgentName.c_str(), remoteAgentName.c_str());
     std::scoped_lock lock(mConnectionsMutex);
     auto it = mConnections.find(remoteAgentName);
     if (it != mConnections.end())
@@ -715,7 +768,7 @@ AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentN
         {
             m_Agent->invalidateRemoteAgent(remoteAgentName);
             it->second->setHasLoadRemoteAgent(true);
-            TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "set has load remote agent to true");
+            TLLM_LOG_DEBUG(mRank, "set has load remote agent to true");
             m_Agent->loadRemoteAgent(remoteAgentName, AgentDesc{metadata.value()});
         }
         return it->second.get();
@@ -725,16 +778,16 @@ AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentN
     {
         if (metadata.has_value())
         {
-            TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "mAgentName: %s connect to %s with loadRemoteAgent",
-                mAgentName.c_str(), remoteAgentName.c_str());
+            TLLM_LOG_DEBUG(mRank, "mAgentName: %s connect to %s with loadRemoteAgent", mAgentName.c_str(),
+                remoteAgentName.c_str());
             m_Agent->loadRemoteAgent(remoteAgentName, AgentDesc{metadata.value()});
             hasLoadRemoteAgent = true;
         }
         else
         {
             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());
+            TLLM_LOG_DEBUG(mRank, "mAgentName: %s connect to %s with loadRemoteAgent", mAgentName.c_str(),
+                remoteAgentName.c_str());
             m_Agent->loadRemoteAgent(remoteAgentName, connectionInfo);
         }
     }

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

@@ -345,6 +345,8 @@ class AgentConnectionManager : public ConnectionManager
     std::unordered_map<std::string, std::list<std::string>> mUnhandledNotifications;
     std::unique_ptr<BaseTransferAgent> m_Agent;
     int mDeviceId;
+    int mRank{0};
+    int mWorldSize{1};
     std::string mAgentName;
     MemoryDescs mRegMemDescs;
     std::atomic<bool> mIsRunning{true};

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION &
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp

@@ -198,20 +198,24 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
             "__init__",
             [](kvc::BaseAgentConfig* self, std::string name, bool use_prog_thread, bool multi_thread,
                 bool use_listen_thread, bool enable_telemetry,
-                std::unordered_map<std::string, std::string> backend_params)
+                std::unordered_map<std::string, std::string> backend_params, std::optional<int> rank,
+                std::optional<int> world_size)
             {
                 new (self) kvc::BaseAgentConfig{std::move(name), use_prog_thread, multi_thread, use_listen_thread,
-                    enable_telemetry, std::move(backend_params)};
+                    enable_telemetry, std::move(backend_params), rank, world_size};
             },
             nb::arg("name"), nb::arg("use_prog_thread") = true, nb::arg("multi_thread") = false,
             nb::arg("use_listen_thread") = false, nb::arg("enable_telemetry") = false,
-            nb::arg("backend_params") = std::unordered_map<std::string, std::string>{})
+            nb::arg("backend_params") = std::unordered_map<std::string, std::string>{}, nb::arg("rank") = std::nullopt,
+            nb::arg("world_size") = std::nullopt)
         .def_rw("name", &kvc::BaseAgentConfig::mName)
         .def_rw("use_prog_thread", &kvc::BaseAgentConfig::useProgThread)
         .def_rw("multi_thread", &kvc::BaseAgentConfig::multiThread)
         .def_rw("use_listen_thread", &kvc::BaseAgentConfig::useListenThread)
         .def_rw("enable_telemetry", &kvc::BaseAgentConfig::enableTelemetry)
-        .def_rw("backend_params", &kvc::BaseAgentConfig::backendParams);
+        .def_rw("backend_params", &kvc::BaseAgentConfig::backendParams)
+        .def_rw("rank", &kvc::BaseAgentConfig::rank)
+        .def_rw("world_size", &kvc::BaseAgentConfig::worldSize);
 
     // BaseTransferAgent class (abstract base)
     // All transfer-engine operations release the GIL: they may block on NIXL /