[TRTLLM-12728][feat] Add WideEP FT E2E MVP integration prototype (DO NOT REVIEW)

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,23 @@ enum class MpiTag : int
 
     // KvCacheEventManager
     kKvCacheEventSize = 1026,
-    kKvCacheEvent = 1027
+    kKvCacheEvent = 1027,
+
+    // Fault-tolerant NCCL communicator rendezvous. Each ownership path and
+    // protocol phase has a distinct tag so stale or simultaneous messages
+    // cannot be decoded as another phase's payload.
+    kNcclCommReady = 1028,
+    kNcclCommUniqueId = 1029,
+    kNcclCommAck = 1030,
+    kNcclPpReady = 1031,
+    kNcclPpUniqueId = 1032,
+    kNcclPpAck = 1033,
+
+    // Seeds used with the canonical group to derive MPI_Comm_create_group
+    // tags for dedicated, pre-failure NCCL rendezvous channels. These are
+    // separate from message tags and distinguish raw-op and PP ownership.
+    kNcclCommControl = 1034,
+    kNcclPpControl = 1035
 };
 
 } // namespace tensorrt_llm::mpi

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

@@ -0,0 +1,52 @@
+/*
+ * 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 <mutex>
+
+namespace tensorrt_llm::runtime
+{
+
+// NCCL host APIs must not run concurrently from different threads for
+// communicators associated with the same CUDA device. TRT-LLM can own both a
+// PP communicator and several cached raw-op communicators, each with its own
+// watchdog, so a per-communicator mutex is insufficient. A process-wide gate
+// is deliberately stricter than a per-device gate and avoids relying on CUDA
+// thread-local device state in watchdog threads.
+//
+// Recursive locking lets low-level error paths abort a communicator while the
+// operation wrapper already owns the gate. A grouped operation keeps one lock
+// alive from ncclGroupStart through ncclGroupEnd.
+inline std::recursive_mutex& getNcclHostApiMutex()
+{
+    // Communicator registries and watchdogs are also process-lifetime
+    // singletons, with cross-translation-unit destruction order unspecified.
+    // Keep the gate alive until process exit so late communicator teardown can
+    // never touch an already-destroyed mutex.
+    static auto* mutex = new std::recursive_mutex;
+    return *mutex;
+}
+
+using NcclHostApiLock = std::unique_lock<std::recursive_mutex>;
+
+inline NcclHostApiLock acquireNcclHostApiLock()
+{
+    return NcclHostApiLock{getNcclHostApiMutex()};
+}
+
+} // namespace tensorrt_llm::runtime

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

@@ -0,0 +1,102 @@
+/*
+ * 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 "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/runtime/utils/mpiUtils.h"
+
+#include <chrono>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#if ENABLE_MULTI_DEVICE
+#include <nccl.h>
+#endif
+
+namespace tensorrt_llm::runtime
+{
+
+#if ENABLE_MULTI_DEVICE
+
+struct NcclUniqueIdRendezvousTags
+{
+    mpi::MpiTag ready;
+    mpi::MpiTag id;
+    mpi::MpiTag ack;
+};
+
+//! An MPI control channel dedicated to NCCL unique-ID rendezvous.
+//!
+//! The channel is created collectively by the initial NCCL group while every
+//! member is healthy. It retains a process-lifetime MPI communicator whose
+//! error handler is MPI_ERRORS_RETURN, leaving the parent/session communicator unchanged. The
+//! original parent-communicator ranks are retained so a later survivor subset
+//! can keep using stable world-rank IDs even though MPI ranks in this channel
+//! are compact.
+//!
+//! This isolates rendezvous traffic and error handling, but it is not a ULFM
+//! communicator repair. Post-failure point-to-point progress still requires an
+//! MPI implementation and launcher configured to let survivors continue.
+class NcclUniqueIdRendezvousComm
+{
+public:
+    NcclUniqueIdRendezvousComm(mpi::MpiComm comm, std::vector<int> worldRanks, int worldRank);
+
+    NcclUniqueIdRendezvousComm(NcclUniqueIdRendezvousComm const&) = delete;
+    NcclUniqueIdRendezvousComm& operator=(NcclUniqueIdRendezvousComm const&) = delete;
+
+    [[nodiscard]] mpi::MpiComm const& mpiComm() const noexcept;
+    [[nodiscard]] std::vector<int> const& worldRanks() const noexcept;
+    [[nodiscard]] int worldRank() const noexcept;
+    [[nodiscard]] int commRank(int worldRank) const;
+    [[nodiscard]] int worldRank(int commRank) const;
+
+private:
+    mpi::MpiComm mComm;
+    std::vector<int> mWorldRanks;
+    int mWorldRank;
+};
+
+//! Create a dedicated control channel over initialRanks.
+//!
+//! Every member of initialRanks must call this before any process failure.
+//! Non-members must not call. creationTagSeed and the canonical group derive a
+//! bounded, deterministic MPI creation tag so concurrently-created overlapping
+//! groups do not normally cross-pair.
+std::shared_ptr<NcclUniqueIdRendezvousComm> createNcclUniqueIdRendezvousComm(
+    std::vector<int> const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed);
+
+// Exchange a fresh NCCL unique ID among only activeRanks. READY/ID/ACK tokens
+// prevent delayed eager MPI messages from an earlier timed-out attempt from
+// pairing different IDs across ranks. rendezvousId is a coordinator-provided
+// logical attempt identity: all ranks in one attempt must pass the same value,
+// and retries for the same logical communicator must use strictly increasing
+// values. The protocol discards older same-communicator messages after seeing
+// a later value while retaining messages for future values and other groups.
+// All ranks use original/global rank IDs.
+// Callers must keep one logical communicator per (control channel, tag set,
+// active-rank set), or construct multiple instances in the same order on every
+// rank. Recovery reuses the pre-failure control channel and performs no MPI
+// collective.
+ncclUniqueId exchangeNcclUniqueId(std::vector<int> const& activeRanks, NcclUniqueIdRendezvousComm const& controlComm,
+    NcclUniqueIdRendezvousTags tags, std::uint64_t rendezvousId, std::chrono::steady_clock::time_point deadline);
+
+#endif // ENABLE_MULTI_DEVICE
+
+} // namespace tensorrt_llm::runtime

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -33,6 +33,7 @@
 #include "tensorrt_llm/runtime/utils/debugUtils.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 
 using namespace tensorrt_llm::kernels;
@@ -363,20 +364,26 @@ int AttentionOp::ulyssesContextPreprocess(T const* input, T* output, T* buffer,
 
     // Do all to all
 #if ENABLE_MULTI_DEVICE
-    ncclGroupStart();
-    for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
     {
-        if (cpIdx != mCpRank)
+        auto commLease = acquireComm(mCpNcclComm);
+        auto const comm = commLease.get();
+        auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context preprocess)");
+        commLease.groupStart("ncclGroupStart(ulysses context preprocess)");
+        auto checkNcclEnqueue = [&](ncclResult_t result)
+        { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context preprocess)"); };
+        for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
         {
-            NCCLCHECK(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
-                (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm,
-                stream));
-            NCCLCHECK(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
-                (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm,
-                stream));
+            if (cpIdx != mCpRank)
+            {
+                checkNcclEnqueue(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
+                    (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream));
+                checkNcclEnqueue(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
+                    (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream));
+            }
         }
+        commLease.groupEnd("ncclGroupEnd(ulysses context preprocess)");
+        commLease.track(watchdogToken, stream);
     }
-    ncclGroupEnd();
     sync_check_cuda_error(stream);
 #endif // ENABLE_MULTI_DEVICE
 
@@ -422,28 +429,36 @@ int AttentionOp::ulyssesContextPostprocess(T* input, T* output, T* buffer, Enque
     // all-to-all
 #if ENABLE_MULTI_DEVICE
     size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads;
-    ncclGroupStart();
-    for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
     {
-        if (cpIdx != mCpRank)
+        auto commLease = acquireComm(mCpNcclComm);
+        auto const comm = commLease.get();
+        auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context postprocess)");
+        commLease.groupStart("ncclGroupStart(ulysses context postprocess)");
+        auto checkNcclEnqueue = [&](ncclResult_t result)
+        { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context postprocess)"); };
+        for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
         {
-            if (mFP8AttenOutput)
-            {
-                NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8,
-                    cpIdx, *mCpNcclComm, stream));
-                NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8,
-                    cpIdx, *mCpNcclComm, stream));
-            }
-            else
+            if (cpIdx != mCpRank)
             {
-                NCCLCHECK(ncclSend(
-                    buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream));
-                NCCLCHECK(ncclRecv(
-                    input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream));
+                if (mFP8AttenOutput)
+                {
+                    checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum,
+                        ncclInt8, cpIdx, comm, stream));
+                    checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum,
+                        ncclInt8, cpIdx, comm, stream));
+                }
+                else
+                {
+                    checkNcclEnqueue(ncclSend(
+                        buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream));
+                    checkNcclEnqueue(
+                        ncclRecv(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream));
+                }
             }
         }
+        commLease.groupEnd("ncclGroupEnd(ulysses context postprocess)");
+        commLease.track(watchdogToken, stream);
     }
-    ncclGroupEnd();
 #endif // ENABLE_MULTI_DEVICE
 
     // transpose_1_reverse + view
@@ -489,20 +504,26 @@ int AttentionOp::ulyssesGenerationPreprocess(
 #if ENABLE_MULTI_DEVICE
     auto const partialHeads = mNumAttnHeads + 2 * mNumAttnKVHeads;
 
-    ncclGroupStart();
-    for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
     {
-        if (cpIdx != mCpRank)
+        auto commLease = acquireComm(mCpNcclComm);
+        auto const comm = commLease.get();
+        auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation preprocess)");
+        commLease.groupStart("ncclGroupStart(ulysses generation preprocess)");
+        auto checkNcclEnqueue = [&](ncclResult_t result)
+        { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation preprocess)"); };
+        for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
         {
-            NCCLCHECK(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
-                (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm,
-                stream));
-            NCCLCHECK(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
-                (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm,
-                stream));
+            if (cpIdx != mCpRank)
+            {
+                checkNcclEnqueue(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
+                    (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream));
+                checkNcclEnqueue(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads),
+                    (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream));
+            }
         }
+        commLease.groupEnd("ncclGroupEnd(ulysses generation preprocess)");
+        commLease.track(watchdogToken, stream);
     }
-    ncclGroupEnd();
     sync_check_cuda_error(stream);
 #endif // ENABLE_MULTI_DEVICE
     return 0;
@@ -525,28 +546,36 @@ int AttentionOp::ulyssesGenerationPostprocess(T* input, T* output, T* buffer, in
     // do all-to-all
 #if ENABLE_MULTI_DEVICE
     size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads;
-    ncclGroupStart();
-    for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
     {
-        if (cpIdx != mCpRank)
+        auto commLease = acquireComm(mCpNcclComm);
+        auto const comm = commLease.get();
+        auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation postprocess)");
+        commLease.groupStart("ncclGroupStart(ulysses generation postprocess)");
+        auto checkNcclEnqueue = [&](ncclResult_t result)
+        { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation postprocess)"); };
+        for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++)
         {
-            if (mFP8AttenOutput)
+            if (cpIdx != mCpRank)
             {
-                NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8,
-                    cpIdx, *mCpNcclComm, stream));
-                NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8,
-                    cpIdx, *mCpNcclComm, stream));
-            }
-            else
-            {
-                NCCLCHECK(ncclSend(
-                    input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream));
-                NCCLCHECK(ncclRecv(
-                    buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream));
+                if (mFP8AttenOutput)
+                {
+                    checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum,
+                        ncclInt8, cpIdx, comm, stream));
+                    checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum,
+                        ncclInt8, cpIdx, comm, stream));
+                }
+                else
+                {
+                    checkNcclEnqueue(
+                        ncclSend(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream));
+                    checkNcclEnqueue(ncclRecv(
+                        buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream));
+                }
             }
         }
+        commLease.groupEnd("ncclGroupEnd(ulysses generation postprocess)");
+        commLease.track(watchdogToken, stream);
     }
-    ncclGroupEnd();
 #endif // ENABLE_MULTI_DEVICE
 
     // do transpose_1_reverse
@@ -3198,11 +3227,13 @@ int AttentionOp::initialize() noexcept
         return 0;
     }
 #if ENABLE_MULTI_DEVICE
-    if (mCpSize > 1 && COMM_SESSION.getSize() > 1)
+    if (mCpSize > 1)
     {
-        TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank());
+        TLLM_CHECK_WITH_INFO(mCpGroup.size() == static_cast<std::size_t>(mCpSize),
+            "Context-parallel group size %zu does not match configured CP size %d", mCpGroup.size(), mCpSize);
+        TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__);
         mCpNcclComm = getComm(mCpGroup);
-        TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank());
+        TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__);
     }
 #endif // ENABLE_MULTI_DEVICE
     return 0;