[None][feat] Use pooled NCCL windows for expert-parallel dispatch and combine

cpp/tensorrt_llm/common/ncclUtils.cpp

@@ -57,20 +57,33 @@ struct CudaMallocGuard
     CudaMallocGuard& operator=(CudaMallocGuard const&) = delete;
 };
 
-// RAII guard for ncclMemAlloc. Frees the pointer on destruction, logging a warning on failure.
+// Roll back a window before freeing its allocation; failed deregistration retains the storage.
 struct NcclMemGuard
 {
     void* ptr{nullptr};
+    ncclComm_t comm;
+    ncclWindow_t window{nullptr};
 
-    explicit NcclMemGuard(void* p) noexcept
+    NcclMemGuard(void* p, ncclComm_t c) noexcept
         : ptr(p)
+        , comm(c)
     {
     }
 
     ~NcclMemGuard()
     {
         if (ptr)
         {
+            if (window)
+            {
+                auto const result = ncclCommWindowDeregister(comm, window);
+                TLLM_NCCL_CHECK_WARN(result);
+                if (result != ncclSuccess)
+                {
+                    TLLM_LOG_WARNING("[NCCLUtil] Retaining allocation after failed window rollback.");
+                    return;
+                }
+            }
             TLLM_NCCL_CHECK_WARN(ncclMemFree(ptr));
         }
     }
@@ -981,12 +994,18 @@ void NCCLWindowAllocator::quarantineTensorLeaseLocked(TensorLease const& tensorL
             entry.graphPoolOwner.reset();
             entry.releaseOwnerWhenUnused = false;
         }
-        TLLM_LOG_DEBUG("[NCCLUtil] Quarantining NCCL window lease %llu on comm %p after a failed tensor lease scope.",
+        TLLM_LOG_DEBUG("[NCCLUtil] Quarantining NCCL window lease %llu on comm %p.",
             static_cast<unsigned long long>(tensorLease.leaseId), static_cast<void*>(tensorLease.comm));
         return;
     }
 }
 
+void NCCLWindowAllocator::abandonReservation(ncclComm_t comm, void* ptr, uint64_t leaseId)
+{
+    std::lock_guard<std::mutex> lock(mMutex);
+    quarantineTensorLeaseLocked(TensorLease{comm, ptr, leaseId});
+}
+
 void NCCLWindowAllocator::releaseBufferFromDestructor(ncclComm_t comm, void* ptr, uint64_t leaseId)
 {
     if (!comm || !ptr || leaseId == 0)
@@ -1123,7 +1142,7 @@ bool NCCLWindowAllocator::isCommValid(ncclComm_t comm) const noexcept
 }
 
 NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(
-    ncclComm_t comm, size_t size, int handle, cudaStream_t stream)
+    ncclComm_t comm, size_t size, int handle, cudaStream_t stream, WindowRegisterFunc registerWindow)
 {
     // Step 1: Pre-allocate the rank-sync flag before ncclMemAlloc. ncclMemAlloc can fail
     // asymmetrically with ncclUnhandledCudaError on configurations where the symmetric/VMM path
@@ -1143,7 +1162,7 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(
     void* ncclPtr = nullptr;
     TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size));
     int const localAllocOk = (ncclPtr != nullptr) ? 1 : 0;
-    NcclMemGuard ncclGuard{ncclPtr}; // frees ncclPtr on any early return or exception
+    NcclMemGuard ncclGuard{ncclPtr, comm}; // frees ncclPtr on any early return or exception
     clearCudaErrorIfSymmetricAllocationFailed(localAllocOk);
 
     // Step 3: ncclCommWindowRegister is collective. If any rank skips it, all other ranks hang.
@@ -1176,12 +1195,30 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(
 
     // Step 4: Register with NCCL as a window. This is collective, so all ranks must reach it.
     // Failure here is non-fatal: warn and fall back to regular allreduce.
-    // ncclGuard frees ncclPtr on return.
+    // Initialize the vote before registration can leave a CUDA error behind.
+    TLLM_CUDA_CHECK(cudaMemsetAsync(rankSyncFlag, 0, sizeof(int), stream));
     ncclWindow_t window = nullptr;
-    ncclResult_t const regResult = ncclCommWindowRegister(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC);
+    ncclResult_t const regResult = registerWindow(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC);
+    ncclGuard.window = window;
     TLLM_NCCL_CHECK_WARN(regResult);
-    if (regResult != ncclSuccess)
+    int const localRegisterOk = (regResult == ncclSuccess && window != nullptr) ? 1 : 0;
+    // A collective registration call does not establish identical local outcomes.
+    // Agree before adding any buffer to the pool, including success with a null window.
+    clearCudaErrorIfSymmetricAllocationFailed(localRegisterOk);
+    if (localRegisterOk)
     {
+        TLLM_CUDA_CHECK_WARN(
+            cudaMemcpyAsync(rankSyncFlag, &localRegisterOk, sizeof(int), cudaMemcpyHostToDevice, stream));
+    }
+    TLLM_NCCL_CHECK(ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream));
+    TLLM_CUDA_CHECK(cudaStreamSynchronize(stream));
+    int allRegisterOk = 0;
+    TLLM_CUDA_CHECK(cudaMemcpy(&allRegisterOk, rankSyncFlag, sizeof(int), cudaMemcpyDeviceToHost));
+    if (!allRegisterOk)
+    {
+        // No caller has seen these windows. Local deregistration rolls back the
+        // successful registrations before their allocations are freed by ncclGuard.
+        TLLM_LOG_WARNING("[NCCLUtil] Window registration unavailable on at least one rank; falling back on all ranks.");
         return NCCLWindowBuffer{};
     }
 

cpp/tensorrt_llm/common/ncclUtils.h

@@ -290,6 +290,9 @@ class NCCLWindowAllocator
     // teardown because destructor timing and stream identity are not rank-deterministic.
     void releaseBufferFromDestructor(ncclComm_t comm, void* ptr, uint64_t leaseId);
 
+    // Explicitly quarantine a raw reservation until communicator teardown. Never makes it reusable.
+    void abandonReservation(ncclComm_t comm, void* ptr, uint64_t leaseId);
+
     // Get the window handle for a specific buffer pointer
     ncclWindow_t getWindow(ncclComm_t comm, void* ptr) const;
 
@@ -322,7 +325,9 @@ class NCCLWindowAllocator
     ~NCCLWindowAllocator() = default;
 
     // Allocate a new buffer and register it with NCCL as a window
-    NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, cudaStream_t stream);
+    using WindowRegisterFunc = ncclResult_t (*)(ncclComm_t, void*, size_t, ncclWindow_t*, int);
+    NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, cudaStream_t stream,
+        WindowRegisterFunc registerWindow = ncclCommWindowRegister);
 
     // Record a failed new symmetric allocation (assumes mMutex is already locked).
     void recordSymmetricFailureLocked(ncclComm_t comm, size_t size);

cpp/tensorrt_llm/thop/allocateOutput.cpp

@@ -28,6 +28,8 @@
 
 #include "tensorrt_llm/thop/outputTensor.h"
 
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/custom_class.h>
 #include <torch/extension.h>
 
 #include <vector>
@@ -37,6 +39,102 @@ TRTLLM_NAMESPACE_BEGIN
 namespace torch_ext
 {
 
+//! A context-owned checkout, independent of operation-scoped tensor leases.
+//! Construction and release must occur in identical rank order outside capture.
+//! Before use, callers must drain prior users on every rank. Before release,
+//! callers must drain all users on every rank and retire graphs using this tensor.
+class NcclWindowReservation : public torch::CustomClassHolder
+{
+public:
+    NcclWindowReservation(at::Tensor like, std::vector<int64_t> group)
+        : mTensor(like)
+    {
+        TORCH_CHECK(like.is_cuda(), "NCCL window reservations require a CUDA tensor");
+        c10::cuda::CUDAGuard const guard(like.device());
+        cudaStreamCaptureStatus status{};
+        TLLM_CUDA_CHECK(cudaStreamIsCapturing(at::cuda::getCurrentCUDAStream().stream(), &status));
+        TORCH_CHECK(status == cudaStreamCaptureStatusNone, "Reserve NCCL windows before CUDA graph capture");
+#if ENABLE_MULTI_DEVICE
+        std::set<int> const ranks(group.begin(), group.end());
+        mComm = getComm(ranks);
+        if (!mComm || !*mComm)
+        {
+            return;
+        }
+        mBuffer = common::nccl_util::NCCLWindowAllocator::getInstance().requestBuffer(
+            *mComm, like.nbytes(), at::cuda::getCurrentCUDAStream().stream());
+        if (mBuffer.isValid())
+        {
+            // Deliberately do not register a temporary tensor lease. The reservation
+            // owns the checkout; tensor views cannot transfer it to a module scope.
+            mTensor = torch::from_blob(mBuffer.ptr, like.sizes(), [comm = mComm](void*){}, like.options());
+        }
+#endif
+    }
+
+    ~NcclWindowReservation() override
+    {
+#if ENABLE_MULTI_DEVICE
+        if (mBuffer.isValid())
+        {
+            // Missing explicit release must never make storage eligible via GC timing.
+            common::nccl_util::NCCLWindowAllocator::getInstance().releaseBufferFromDestructor(
+                *mComm, mBuffer.ptr, mBuffer.leaseId);
+        }
+#endif
+    }
+
+    //! Quarantine without CUDA or collective calls; retain registration until communicator teardown.
+    void abandon()
+    {
+#if ENABLE_MULTI_DEVICE
+        if (mBuffer.isValid())
+        {
+            common::nccl_util::NCCLWindowAllocator::getInstance().abandonReservation(
+                *mComm, mBuffer.ptr, mBuffer.leaseId);
+            mBuffer = {};
+        }
+#endif
+    }
+
+    at::Tensor getTensor() const
+    {
+        return mTensor;
+    }
+
+    //! Return the raw ncclWindow_t encoded as an integer, or 0 on a reservation miss.
+    int64_t getWindow() const
+    {
+#if ENABLE_MULTI_DEVICE
+        return static_cast<int64_t>(reinterpret_cast<uintptr_t>(mBuffer.window));
+#else
+        return 0;
+#endif
+    }
+
+    //! Release after the caller's collective quiescence boundary; idempotent.
+    void release()
+    {
+#if ENABLE_MULTI_DEVICE
+        if (mBuffer.isValid())
+        {
+            c10::cuda::CUDAGuard const guard(mTensor.device());
+            common::nccl_util::NCCLWindowAllocator::getInstance().releaseBuffer(
+                *mComm, mBuffer.ptr, at::cuda::getCurrentCUDAStream().stream());
+            mBuffer = {};
+        }
+#endif
+        mTensor = at::Tensor();
+    }
+
+private:
+    at::Tensor mTensor;
+#if ENABLE_MULTI_DEVICE
+    std::shared_ptr<ncclComm_t> mComm;
+    common::nccl_util::NCCLWindowBuffer mBuffer{};
+#endif
+};
+
 std::tuple<at::Tensor, int64_t> allocateOutputOp(at::Tensor const& like, int64_t output_buffer_kind,
     c10::optional<torch::List<int64_t>> group, c10::optional<at::IntArrayRef> shape,
     c10::optional<at::ScalarType> out_dtype)
@@ -93,6 +191,12 @@ TRTLLM_NAMESPACE_END
 
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
+    m.class_<tensorrt_llm::torch_ext::NcclWindowReservation>("NCCLWindowReservation")
+        .def(torch::init<at::Tensor, std::vector<int64_t>>())
+        .def("tensor", &tensorrt_llm::torch_ext::NcclWindowReservation::getTensor)
+        .def("window", &tensorrt_llm::torch_ext::NcclWindowReservation::getWindow)
+        .def("abandon", &tensorrt_llm::torch_ext::NcclWindowReservation::abandon)
+        .def("release", &tensorrt_llm::torch_ext::NcclWindowReservation::release);
     m.def(
         "allocate_output(Tensor like, int output_buffer_kind, int[]? group, "
         "int[]? shape=None, ScalarType? out_dtype=None) -> (Tensor, int)");

cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp

@@ -48,6 +48,12 @@ namespace tensorrt_llm::common::nccl_util
 class NCCLWindowAllocatorTestAccess
 {
 public:
+    static NCCLWindowBuffer allocateWithRegistration(
+        NCCLWindowAllocator& allocator, ncclComm_t comm, NCCLWindowAllocator::WindowRegisterFunc registerWindow)
+    {
+        return allocator.allocateAndRegisterBuffer(comm, 1024 * 1024, 0, nullptr, registerWindow);
+    }
+
     static void recordSymmetricFailure(NCCLWindowAllocator& allocator, ncclComm_t comm, size_t size)
     {
         std::lock_guard<std::mutex> lock(allocator.mMutex);
@@ -467,6 +473,52 @@ TEST_F(NCCLWindowAllocatorTest, BasicAllocation)
     allocator.releaseBuffer(*mComm, buffer.ptr);
 }
 
+TEST_F(NCCLWindowAllocatorTest, AsymmetricRegistrationFallsBackOnEveryRank)
+{
+    auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();
+    // Both ranks enter real registration; only rank zero reports an unavailable window.
+    // Exercise both an explicit error and NCCL's successful-but-null result.
+    for (bool const nullSuccess : {false, true})
+    {
+        auto const registerWithError = [](ncclComm_t comm, void* ptr, size_t size, ncclWindow_t* window, int flags)
+        {
+            auto const result = ncclCommWindowRegister(comm, ptr, size, window, flags);
+            if (result == ncclSuccess && mpi::MpiComm::world().getRank() == 0)
+            {
+                TLLM_NCCL_CHECK(ncclCommWindowDeregister(comm, *window));
+                *window = nullptr;
+                return ncclSystemError;
+            }
+            return result;
+        };
+        auto const registerWithNull = [](ncclComm_t comm, void* ptr, size_t size, ncclWindow_t* window, int flags)
+        {
+            auto const result = ncclCommWindowRegister(comm, ptr, size, window, flags);
+            if (result == ncclSuccess && mpi::MpiComm::world().getRank() == 0)
+            {
+                TLLM_NCCL_CHECK(ncclCommWindowDeregister(comm, *window));
+                *window = nullptr;
+            }
+            return result;
+        };
+        auto const buffer = nccl_util::NCCLWindowAllocatorTestAccess::allocateWithRegistration(
+            allocator, *mComm, nullSuccess ? +registerWithNull : +registerWithError);
+        EXPECT_FALSE(buffer.isValid());
+        if (buffer.isValid())
+        {
+            TLLM_NCCL_CHECK(ncclCommWindowDeregister(*mComm, buffer.window));
+            TLLM_NCCL_CHECK(ncclMemFree(buffer.ptr));
+        }
+    }
+    // Rollback must leave the communicator usable for a subsequent pooled allocation.
+    auto const buffer = allocator.requestBuffer(*mComm, 1024 * 1024);
+    EXPECT_TRUE(buffer.isValid());
+    if (buffer.isValid())
+    {
+        allocator.releaseBuffer(*mComm, buffer.ptr);
+    }
+}
+
 TEST_F(NCCLWindowAllocatorTest, ExplicitStreamAllocationAndCapture)
 {
     auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();