[pull] main from NVIDIA:main

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -27,6 +27,7 @@
 #include "tensorrt_llm/executor/dataTransceiverState.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include "tensorrt_llm/runtime/utils/pgUtils.h"
+#include <cstddef>
 #include <future>
 #include <memory>
 #include <mutex>
@@ -55,6 +56,7 @@ class BaseKVCacheManager;
 
 class CacheSender;
 class CacheReceiver;
+class ContextTransferCoordinator;
 
 class CacheTransceiverComm
 {
@@ -149,6 +151,25 @@ class CacheTransceiverComm
         TLLM_THROW("Input arguments only supported in mpi");
     }
 
+    [[nodiscard]] std::unique_ptr<mpi::MpiRequest> sendAsync(
+        void const* buffer, std::size_t size, mpi::MpiType dtype, int dest, mpi::MpiTag tag) const
+    {
+        TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI.");
+        return mMpiComm->sendAsync(buffer, size, dtype, dest, tag);
+    }
+
+    [[nodiscard]] bool iprobe(int source, mpi::MpiTag tag, MPI_Status* status) const
+    {
+        TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI.");
+        return mMpiComm->iprobe(source, tag, status);
+    }
+
+    void recv(void* buffer, std::size_t size, mpi::MpiType dtype, int source, mpi::MpiTag tag) const
+    {
+        TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI.");
+        static_cast<void>(mMpiComm->recv(buffer, size, dtype, source, tag));
+    }
+
     CacheTransceiverComm split(int color, int key)
     {
         if (isMpi())
@@ -307,6 +328,7 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     std::shared_ptr<CacheTransceiverComm> mGroupComm;
     std::shared_ptr<CacheTransceiverComm> mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm;
+    std::unique_ptr<ContextTransferCoordinator> mContextTransferCoordinator;
 
     executor::kv_cache::CommState const* mCommState;
     std::unique_ptr<executor::kv_cache::CacheState> mCacheState;

cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h

@@ -0,0 +1,109 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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 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 <cstddef>
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace tensorrt_llm::batch_manager
+{
+
+class CacheTransceiverComm;
+
+enum class ContextTransferVote : std::uint64_t
+{
+    kCompleted = 1,
+    kFailed = 2,
+};
+
+struct ContextTransferConsensusResult
+{
+    std::unordered_set<std::uint64_t> completedRequestIds;
+    std::unordered_set<std::uint64_t> failedRequestIds;
+    std::unordered_set<std::uint64_t> timedOutRequestIds;
+};
+
+//! Accumulates one immutable terminal vote per participant and request.
+class ContextTransferVoteReducer
+{
+public:
+    explicit ContextTransferVoteReducer(int participantCount);
+
+    void recordVote(int participantRank, std::uint64_t requestId, ContextTransferVote vote);
+
+    //! Record a sticky nonterminal timeout proposal. The final outcome still waits for every terminal vote.
+    void recordTimeout(std::uint64_t requestId);
+
+    [[nodiscard]] ContextTransferConsensusResult takeReady();
+
+    void clear() noexcept;
+
+private:
+    struct RequestVotes
+    {
+        explicit RequestVotes(int participantCount)
+            : votes(static_cast<std::size_t>(participantCount), 0)
+        {
+        }
+
+        std::vector<std::uint64_t> votes;
+        int terminalCount{0};
+        bool failed{false};
+        bool timedOut{false};
+        bool timeoutPending{false};
+    };
+
+    int mParticipantCount;
+    std::unordered_map<std::uint64_t, RequestVotes> mRequestVotes;
+};
+
+//! Coordinates asynchronous context-transfer outcomes across a rank group.
+//!
+//! Timeout and terminal events share one ordered stream toward the coordinator. Timeout updates and final commits
+//! share another ordered stream toward followers, so a request that times out can never commit success first.
+class ContextTransferCoordinator
+{
+public:
+    explicit ContextTransferCoordinator(std::shared_ptr<CacheTransceiverComm> comm);
+    ~ContextTransferCoordinator();
+
+    ContextTransferCoordinator(ContextTransferCoordinator const&) = delete;
+    ContextTransferCoordinator& operator=(ContextTransferCoordinator const&) = delete;
+
+    //! Publish this rank's immutable local terminal outcome without waiting for peers.
+    void publishLocalOutcome(std::uint64_t requestId, bool failed);
+
+    //! Publish an idempotent, sticky, nonterminal timeout proposal without waiting for peers.
+    void publishTimeout(std::uint64_t requestId);
+
+    //! Make nonblocking protocol progress and return newly committed global outcomes.
+    [[nodiscard]] ContextTransferConsensusResult poll();
+
+    //! Exchange ordered close markers so no active MPI request outlives its backing buffer. Failure aborts closed.
+    void shutdown() noexcept;
+
+private:
+    class Impl;
+    std::unique_ptr<Impl> mImpl;
+};
+
+} // namespace tensorrt_llm::batch_manager

cpp/include/tensorrt_llm/runtime/utils/mpiTags.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2021-2026, 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.
@@ -71,7 +71,11 @@ enum class MpiTag : int
 
     // KvCacheEventManager
     kKvCacheEventSize = 1026,
-    kKvCacheEvent = 1027
+    kKvCacheEvent = 1027,
+
+    // Asynchronous context-transfer coordination.
+    kContextTransferEvent = 1028,
+    kContextTransferUpdate = 1029
 };
 
 } // namespace tensorrt_llm::mpi

cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2021-2026, 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.
@@ -235,6 +235,17 @@ class MpiRequest
 #endif
     }
 
+    [[nodiscard]] bool isCompleted()
+    {
+#if ENABLE_MULTI_DEVICE
+        int completed = 0;
+        TLLM_MPI_CHECK(MPI_Test(&mRequest, &completed, MPI_STATUS_IGNORE));
+        return completed != 0;
+#else
+        TLLM_THROW("Multi device support is disabled.");
+#endif
+    }
+
     void cancel()
     {
 #if ENABLE_MULTI_DEVICE

cpp/tensorrt_llm/batch_manager/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
@@ -30,6 +30,7 @@ set(SRCS
     capacityScheduler.cpp
     createNewDecoderRequests.cpp
     contextProgress.cpp
+    contextTransferCoordinator.cpp
     dataTransceiver.cpp
     decoderBuffers.cpp
     kvCacheManager.cpp