[TRTLLM-15344][feat] cache transceiver nixl bounce buffer

cpp/CMakeLists.txt

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+# SPDX-FileCopyrightText: Copyright (c) 1993-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
@@ -280,8 +280,15 @@ include_directories(${CMAKE_BINARY_DIR}/_deps/nanobind-src/include)
 
 FetchContent_MakeAvailable(cutlass cxxopts flashmla json msa xgrammar)
 
+# cppzmq (header-only) is the bounce v2 control-plane dependency AND the UCX
+# bootstrap dependency. Make it available whenever NIXL (-> bounce) or UCX is
+# enabled, so bounce does NOT require UCX (its real deps are NIXL + zmq,
+# orthogonal to UCX). ucxx stays UCX-only.
+if(NIXL_ROOT OR ENABLE_UCX)
+  FetchContent_MakeAvailable(cppzmq)
+endif()
 if(ENABLE_UCX)
-  FetchContent_MakeAvailable(cppzmq ucxx)
+  FetchContent_MakeAvailable(ucxx)
 endif()
 
 if(NOT NVTX_DISABLE)

cpp/include/tensorrt_llm/executor/executor.h

@@ -1522,20 +1522,25 @@ class CacheTransceiverConfig
     explicit CacheTransceiverConfig(std::optional<BackendType> backendType = std::nullopt,
         std::optional<size_t> maxNumTokens = std::nullopt, std::optional<int> kvTransferTimeoutMs = std::nullopt,
         std::optional<int> kvTransferSenderFutureTimeoutMs = std::nullopt,
-        std::optional<int> kvTransferPollIntervalMs = kDefaultKvTransferPollIntervalMs);
+        std::optional<int> kvTransferPollIntervalMs = kDefaultKvTransferPollIntervalMs, size_t agentBufferSizeMb = 0,
+        std::map<std::string, std::string> agentBounceParams = {});
 
     bool operator==(CacheTransceiverConfig const& other) const;
     void setBackendType(std::optional<BackendType> backendType);
     void setMaxTokensInBuffer(std::optional<size_t> maxTokensInBuffer);
     void setKvTransferTimeoutMs(std::optional<int> kvTransferTimeoutMs);
     void setKvTransferSenderFutureTimeoutMs(std::optional<int> kvTransferSenderFutureTimeoutMs);
     void setKvTransferPollIntervalMs(std::optional<int> kvTransferPollIntervalMs);
+    void setAgentBufferSizeMb(size_t agentBufferSizeMb);
+    void setAgentBounceParams(std::map<std::string, std::string> agentBounceParams);
 
     [[nodiscard]] std::optional<size_t> getMaxTokensInBuffer() const;
     [[nodiscard]] std::optional<BackendType> getBackendType() const;
     [[nodiscard]] std::optional<int> getKvTransferTimeoutMs() const;
     [[nodiscard]] std::optional<int> getKvTransferSenderFutureTimeoutMs() const;
     [[nodiscard]] std::optional<int> getKvTransferPollIntervalMs() const;
+    [[nodiscard]] size_t getAgentBufferSizeMb() const;
+    [[nodiscard]] std::map<std::string, std::string> const& getAgentBounceParams() const;
 
 private:
     std::optional<BackendType> mBackendType;
@@ -1550,6 +1555,15 @@ class CacheTransceiverConfig
     // @brief Bounded wait interval in milliseconds for polling KV transfer progress when active transfers block
     // disaggregated admission.
     std::optional<int> mKvTransferPollIntervalMs;
+    // @brief Size in MiB of the transfer agent's staging-buffer (bounce) arena, currently
+    // implemented by the NIXL agent. 0 (default) disables the fast path; >0 enables it at that
+    // arena capacity (one shared arena for the sender and receiver roles).
+    size_t mAgentBufferSizeMb{0};
+    // @brief Expert tuning knobs for the bounce pipeline, keyed by the TRTLLM_NIXL_BOUNCE_*
+    // environment-variable names without the prefix and the trailing _BYTES, lowercased (e.g.
+    // TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES -> "max_chunk_size"). Precedence: this map >
+    // environment variable > built-in default. Ignored when mAgentBufferSizeMb == 0.
+    std::map<std::string, std::string> mAgentBounceParams;
 };
 
 /// @brief Configuration class for the model executor

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -251,9 +251,8 @@ struct VramRegionMeta
     size_t chunkSize; ///< 0 = cudaMalloc (no split), >0 = VMM chunk size
 };
 
-// `AgentDesc` represents the unique identifier for reading and writing to the agent.
-// By accessing this identifier, the backend can establish the correct connection.
-// It also carries VMM region metadata so that remote agents can split at chunk boundaries.
+// `AgentDesc` carries the backend metadata needed to connect to an agent, plus optional
+// backend-independent metadata used by higher-level transfer paths.
 class AgentDesc final
 {
 public:
@@ -262,9 +261,10 @@ class AgentDesc final
     {
     }
 
-    AgentDesc(std::string backendAgentDesc, std::vector<VramRegionMeta> vramRegions)
+    AgentDesc(std::string backendAgentDesc, std::vector<VramRegionMeta> vramRegions, std::string bounceHandshake = {})
         : mBackendAgentDesc{std::move(backendAgentDesc)}
         , mVramRegions{std::move(vramRegions)}
+        , mBounceHandshake{std::move(bounceHandshake)}
     {
     }
 
@@ -278,7 +278,18 @@ class AgentDesc final
         return mVramRegions;
     }
 
-    /// Serialize the entire AgentDesc (backend blob + VMM regions) into an opaque string.
+    /// Optional NIXL-bounce capability handshake advertised by this agent.
+    /// An opaque blob encoded/decoded by bounce::encodeHandshake/decodeHandshake: wire version,
+    /// control-channel kind + endpoint, and effective region-size limits. It is empty when bounce is
+    /// disabled. A caller must exchange the complete serialized AgentDesc, rather than only
+    /// getBackendAgentDesc(), for the peer to receive this handshake.
+    [[nodiscard]] std::string const& getBounceHandshake() const noexcept
+    {
+        return mBounceHandshake;
+    }
+
+    /// Serialize the entire AgentDesc (backend blob + VMM regions + bounce handshake) into an
+    /// opaque string.
     [[nodiscard]] std::string serialize() const;
 
     /// Deserialize an opaque string back into an AgentDesc.
@@ -287,6 +298,7 @@ class AgentDesc final
 private:
     std::string mBackendAgentDesc;
     std::vector<VramRegionMeta> mVramRegions;
+    std::string mBounceHandshake;
 };
 
 // `TransferOp` is an enumeration that represents the types of transfer operations.
@@ -371,6 +383,13 @@ class TransferStatus
     {
         return false;
     }
+
+    /// Human-readable detail for the most recent terminal state (empty when unavailable). Backends
+    /// override this so a kFAILURE from wait() can carry its cause to the caller.
+    [[nodiscard]] virtual std::string getLastStatusStr() const
+    {
+        return {};
+    }
 };
 
 struct BaseAgentConfig
@@ -383,6 +402,18 @@ struct BaseAgentConfig
     std::unordered_map<std::string, std::string> backendParams;
     std::optional<int> rank;
     std::optional<int> worldSize;
+    /// Size in MiB of the agent's staging-buffer (bounce) arena, currently implemented by the
+    /// NIXL agent. 0 (default) disables the fast path; >0 enables it at that arena capacity.
+    /// Ignored by agents without bounce support (e.g. mooncake).
+    size_t agentBufferSizeMb{0};
+    /// Expert tuning knobs for the bounce pipeline, keyed by the TRTLLM_NIXL_BOUNCE_*
+    /// environment-variable names without the prefix and the trailing _BYTES, lowercased (e.g.
+    /// TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES -> "max_chunk_size").
+    /// Kept SEPARATE from backendParams on purpose: backendParams is forwarded verbatim to the
+    /// backend plugin (e.g. NIXL createBackend), so bounce keys must not leak into it.
+    /// Precedence: this map > environment variable > built-in default. Ignored when
+    /// agentBufferSizeMb == 0.
+    std::unordered_map<std::string, std::string> bounceParams;
 };
 
 class BaseTransferAgent
@@ -416,8 +447,8 @@ class BaseTransferAgent
     /// @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.
+    /// @brief Fetch the backend-specific connection information for the local agent.
+    /// @return The local connection information. Optional AgentDesc metadata is not included.
     virtual ConnectionInfoType getLocalConnectionInfo() = 0;
 
     /// @brief Initiate the transfer by submitting the request.

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -664,15 +664,17 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
         auto rnnState
             = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt;
         mManager = std::make_unique<tensorrt_llm::executor::kv_cache::AgentConnectionManager>(
-            mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState);
+            mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState,
+            mCacheTransceiverConfig->getAgentBufferSizeMb(), mCacheTransceiverConfig->getAgentBounceParams());
         TLLM_LOG_INFO("NIXL Connection Manager created");
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE)
     {
         auto rnnState
             = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt;
         mManager = std::make_unique<tensorrt_llm::executor::kv_cache::AgentConnectionManager>(
-            mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState);
+            mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState,
+            mCacheTransceiverConfig->getAgentBufferSizeMb(), mCacheTransceiverConfig->getAgentBounceParams());
         TLLM_LOG_INFO("MOONCAKE Connection Manager created");
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI)

cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp

@@ -23,11 +23,14 @@ namespace tensorrt_llm::executor
 
 CacheTransceiverConfig::CacheTransceiverConfig(std::optional<BackendType> backendType,
     std::optional<size_t> maxNumTokens, std::optional<int> kvTransferTimeoutMs,
-    std::optional<int> kvTransferSenderFutureTimeoutMs, std::optional<int> kvTransferPollIntervalMs)
+    std::optional<int> kvTransferSenderFutureTimeoutMs, std::optional<int> kvTransferPollIntervalMs,
+    size_t agentBufferSizeMb, std::map<std::string, std::string> agentBounceParams)
     : mBackendType(backendType)
     , mMaxTokensInBuffer(maxNumTokens)
     , mKvTransferTimeoutMs(kvTransferTimeoutMs)
     , mKvTransferSenderFutureTimeoutMs(kvTransferSenderFutureTimeoutMs)
+    , mAgentBufferSizeMb(agentBufferSizeMb)
+    , mAgentBounceParams(std::move(agentBounceParams))
 {
     setKvTransferPollIntervalMs(kvTransferPollIntervalMs);
 }
@@ -37,7 +40,8 @@ bool CacheTransceiverConfig::operator==(CacheTransceiverConfig const& other) con
     return mMaxTokensInBuffer == other.mMaxTokensInBuffer && mBackendType == other.mBackendType
         && mKvTransferTimeoutMs == other.mKvTransferTimeoutMs
         && mKvTransferSenderFutureTimeoutMs == other.mKvTransferSenderFutureTimeoutMs
-        && mKvTransferPollIntervalMs == other.mKvTransferPollIntervalMs;
+        && mKvTransferPollIntervalMs == other.mKvTransferPollIntervalMs
+        && mAgentBufferSizeMb == other.mAgentBufferSizeMb && mAgentBounceParams == other.mAgentBounceParams;
 }
 
 void CacheTransceiverConfig::setBackendType(std::optional<BackendType> backendType)
@@ -101,4 +105,24 @@ std::optional<int> CacheTransceiverConfig::getKvTransferPollIntervalMs() const
 {
     return mKvTransferPollIntervalMs;
 }
+
+void CacheTransceiverConfig::setAgentBufferSizeMb(size_t agentBufferSizeMb)
+{
+    mAgentBufferSizeMb = agentBufferSizeMb;
+}
+
+size_t CacheTransceiverConfig::getAgentBufferSizeMb() const
+{
+    return mAgentBufferSizeMb;
+}
+
+void CacheTransceiverConfig::setAgentBounceParams(std::map<std::string, std::string> agentBounceParams)
+{
+    mAgentBounceParams = std::move(agentBounceParams);
+}
+
+std::map<std::string, std::string> const& CacheTransceiverConfig::getAgentBounceParams() const
+{
+    return mAgentBounceParams;
+}
 } // namespace tensorrt_llm::executor

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

@@ -386,7 +386,8 @@ std::optional<size_t> AgentConnection::getPreAssignedBufferId(uint8_t kind) cons
 
 AgentConnectionManager::AgentConnectionManager(
     std::vector<batch_manager::BaseTransBufferManager*> cacheTransBufferManagers, CacheState cacheState,
-    std::string const& backendType, std::optional<CacheState::RnnCacheState> rnnCacheState)
+    std::string const& backendType, std::optional<CacheState::RnnCacheState> rnnCacheState, size_t agentBufferSizeMb,
+    std::map<std::string, std::string> const& agentBounceParams)
     : mCacheState(std::move(cacheState))
     , mRnnCacheState(std::move(rnnCacheState))
     , mCacheTransBufferManagers(std::move(cacheTransBufferManagers))
@@ -425,6 +426,11 @@ AgentConnectionManager::AgentConnectionManager(
     BaseAgentConfig config{mAgentName, true, false, true};
     config.rank = mRank;
     config.worldSize = mWorldSize;
+    config.agentBufferSizeMb = agentBufferSizeMb;
+    // BaseAgentConfig keeps bounce knobs separate from backendParams (which is forwarded verbatim
+    // to the backend plugin); convert the ordered map from CacheTransceiverConfig here.
+    config.bounceParams
+        = std::unordered_map<std::string, std::string>(agentBounceParams.begin(), agentBounceParams.end());
     m_Agent = makeTransferAgent(backendType, &config);
     TLLM_CHECK(!mCacheTransBufferManagers.empty());
     mBufferKinds.reserve(mCacheTransBufferManagers.size());

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

@@ -305,7 +305,8 @@ class AgentConnectionManager : public ConnectionManager
 public:
     AgentConnectionManager(std::vector<batch_manager::BaseTransBufferManager*> cacheTransBufferManagers,
         CacheState cacheState, std::string const& backendType,
-        std::optional<CacheState::RnnCacheState> rnnCacheState = std::nullopt);
+        std::optional<CacheState::RnnCacheState> rnnCacheState = std::nullopt, size_t agentBufferSizeMb = 0,
+        std::map<std::string, std::string> const& agentBounceParams = {});
     ~AgentConnectionManager();
     AgentConnection* recvConnect(DataContext const& ctx, void* data, size_t size) override;
     [[nodiscard]] std::vector<Connection const*> getConnections(CommState const& state) override;

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

@@ -42,6 +42,42 @@ if(NIXL_ROOT)
   # Link against CUDA runtime (for cudaMemcpy in posix fallback)
   target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE CUDA::cudart)
 
+  # ------------------------------------------------------------------------
+  # Bounce v2 transport. Its real dependencies are NIXL (data plane, gated by
+  # the NIXL_ROOT block above) + zmq (control plane) — NOT UCX. So compile it
+  # whenever zmq is available, INDEPENDENT of ENABLE_UCX. cppzmq (header-only)
+  # is fetched at the top level when NIXL or UCX is on. Activates the
+  # TLLM_BOUNCE_V2 code paths; runtime use is still opt-in via
+  # CacheTransceiverConfig.agent_buffer_size_mb.
+  # ------------------------------------------------------------------------
+  find_package(PkgConfig)
+  if(PKG_CONFIG_FOUND)
+    pkg_check_modules(ZMQ libzmq) # optional: bounce is skipped if libzmq is
+                                  # absent
+  endif()
+  if(ZMQ_FOUND)
+    target_sources(
+      ${NIXL_WRAPPER_TARGET}
+      PRIVATE bounce/BounceTransferPlan.cpp
+              bounce/BuddyAllocator.cpp
+              bounce/CreditScheduler.cpp
+              bounce/BounceMessage.cpp
+              bounce/BounceArena.cpp
+              bounce/ExecPool.cpp
+              bounce/GatherScatterKernel.cu
+              bounce/ZmqControlChannel.cpp
+              bounce/BounceTransport.cpp)
+    target_include_directories(
+      ${NIXL_WRAPPER_TARGET} PRIVATE ${CMAKE_BINARY_DIR}/_deps/cppzmq-src
+                                     ${ZMQ_INCLUDE_DIRS})
+    target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE ${ZMQ_LIBRARIES}
+                                                         ${NIXL_BUILD_LIBRARY})
+    target_compile_definitions(${NIXL_WRAPPER_TARGET} PRIVATE TLLM_BOUNCE_V2=1)
+    # The .cu device symbols must resolve into this shared lib.
+    set_target_properties(${NIXL_WRAPPER_TARGET}
+                          PROPERTIES CUDA_RESOLVE_DEVICE_SYMBOLS ON)
+  endif()
+
   set(NIXL_ENABLED TRUE)
 else()
   set(NIXL_ENABLED FALSE)

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

@@ -189,7 +189,11 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
     // subclass type is not directly registered (e.g., agents created via factory).
     nb::class_<kvc::TransferStatus>(m, "TransferStatus")
         .def("is_completed", &kvc::TransferStatus::isCompleted, nb::call_guard<nb::gil_scoped_release>())
-        .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard<nb::gil_scoped_release>());
+        .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard<nb::gil_scoped_release>())
+        // Failure detail for the last terminal state (empty if unavailable). Named to match the
+        // lookup in BindingsNixlTransferStatus.last_status_str (nixl/_agent_cpp.py), which is what
+        // the Python transceiver's error log reads.
+        .def("get_last_status_str", &kvc::TransferStatus::getLastStatusStr);
 
     // BaseAgentConfig struct
     nb::class_<kvc::BaseAgentConfig>(m, "BaseAgentConfig")
@@ -199,23 +203,28 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
             [](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::optional<int> rank,
-                std::optional<int> world_size)
+                std::optional<int> world_size, size_t agent_buffer_size_mb,
+                std::unordered_map<std::string, std::string> bounce_params)
             {
                 new (self) kvc::BaseAgentConfig{std::move(name), use_prog_thread, multi_thread, use_listen_thread,
-                    enable_telemetry, std::move(backend_params), rank, world_size};
+                    enable_telemetry, std::move(backend_params), rank, world_size, agent_buffer_size_mb,
+                    std::move(bounce_params)};
             },
             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("rank") = std::nullopt,
-            nb::arg("world_size") = std::nullopt)
+            nb::arg("world_size") = std::nullopt, nb::arg("agent_buffer_size_mb") = 0,
+            nb::arg("bounce_params") = std::unordered_map<std::string, std::string>{})
         .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("rank", &kvc::BaseAgentConfig::rank)
-        .def_rw("world_size", &kvc::BaseAgentConfig::worldSize);
+        .def_rw("world_size", &kvc::BaseAgentConfig::worldSize)
+        .def_rw("agent_buffer_size_mb", &kvc::BaseAgentConfig::agentBufferSizeMb)
+        .def_rw("bounce_params", &kvc::BaseAgentConfig::bounceParams);
 
     // BaseTransferAgent class (abstract base)
     // All transfer-engine operations release the GIL: they may block on NIXL /
@@ -293,7 +302,10 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
         .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages,
             nb::call_guard<nb::gil_scoped_release>())
         .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"),
-            nb::call_guard<nb::gil_scoped_release>());
+            nb::call_guard<nb::gil_scoped_release>())
+        // Programmatic bounce v2 observability (deployment checks and tests; no log parsing).
+        .def_prop_ro("bounce_enabled", &kvc::NixlTransferAgent::isBounceEnabled)
+        .def_prop_ro("bounce_submit_count", &kvc::NixlTransferAgent::getBounceSubmitCount);
 #endif
 
     // NOTE: MooncakeTransferAgent/MooncakeTransferStatus class bindings are intentionally