[https://nvbugs/7026168][fix] Skip eager NCCL-window layer scopes until the first window request

cpp/tensorrt_llm/common/ncclUtils.cpp

@@ -614,6 +614,8 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size
             static_cast<void*>(comm), size);
         return NCCLWindowBuffer();
     }
+    // Recorded whatever the outcome below: every rank makes this request.
+    mBufferRequested.store(true);
 
     TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null");
     TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0");
@@ -1033,6 +1035,48 @@ void NCCLWindowAllocator::endTensorLeaseScope(std::vector<c10::StorageImpl const
     }
 }
 
+void NCCLWindowAllocator::releaseTensorLeaseScopeLeases(
+    std::vector<c10::StorageImpl const*> const& keptStorages, int device, cudaStream_t stream)
+{
+    std::unordered_set<c10::StorageImpl const*> const kept(keptStorages.begin(), keptStorages.end());
+
+    std::lock_guard<std::mutex> lock(mMutex);
+    auto const stackIt = mTensorLeaseScopeStacks.find(getTensorLeaseScopeKey(device));
+    if (stackIt == mTensorLeaseScopeStacks.end() || stackIt->second.empty())
+    {
+        static std::once_flag noScopeWarning;
+        std::call_once(noScopeWarning,
+            []
+            {
+                TLLM_LOG_WARNING(
+                    "[NCCLUtil] No NCCL window tensor lease scope is open on this thread, so no lease is released. "
+                    "Window tensors created outside a scope stay leased until their destructors quarantine their "
+                    "buffers. Run eager models inside nccl_window_tensor_scope.");
+            });
+        return;
+    }
+    auto& scope = stackIt->second.back();
+    for (auto storageIt = scope.storages.begin(); storageIt != scope.storages.end();)
+    {
+        if (kept.find(*storageIt) != kept.end())
+        {
+            ++storageIt;
+            continue;
+        }
+        auto const leaseIt = mTensorLeases.find(*storageIt);
+        TLLM_CHECK_WITH_INFO(leaseIt != mTensorLeases.end(), "NCCL window tensor scope contains an unknown storage");
+        auto const releasedLease = leaseIt->second;
+        mTensorLeases.erase(leaseIt);
+        storageIt = scope.storages.erase(storageIt);
+        releaseTensorLeaseLocked(releasedLease, stream);
+    }
+    for (auto const& detachedLease : scope.detachedLeases)
+    {
+        releaseTensorLeaseLocked(detachedLease, stream);
+    }
+    scope.detachedLeases.clear();
+}
+
 bool NCCLWindowAllocator::releaseTensorBuffer(ncclComm_t comm, c10::StorageImpl const* storage, cudaStream_t stream)
 {
     if (!comm || !storage)
@@ -1238,6 +1282,11 @@ size_t NCCLWindowAllocator::getBufferCount(ncclComm_t comm) const
     return commIt != mBufferPool.end() ? commIt->second.size() : 0;
 }
 
+bool NCCLWindowAllocator::hasRequestedBuffers() const noexcept
+{
+    return mBufferRequested.load();
+}
+
 size_t NCCLWindowAllocator::getBufferInUseCount(ncclComm_t comm) const
 {
     std::lock_guard<std::mutex> lock(mMutex);

cpp/tensorrt_llm/common/ncclUtils.h

@@ -277,6 +277,9 @@ class NCCLWindowAllocator
     // New allocation/registration during capture requires NCCL runtime 2.30.7+ and a graph-pool owner.
     // Uses the supplied stream for capture detection and stream-scoped reuse eligibility.
     // Callers using a non-default stream must pass it explicitly.
+    // Every rank of comm makes each request at the same point: a new allocation is collective, and
+    // the first request also changes hasRequestedBuffers(). A request that only some ranks make must
+    // run with allocation suppressed (setWindowAllocationSuppressed).
     NCCLWindowBuffer requestBuffer(ncclComm_t comm, size_t size, cudaStream_t stream = nullptr);
 
     // Select the allocation domain on this host thread. owner == -1 explicitly selects eager;
@@ -330,6 +333,13 @@ class NCCLWindowAllocator
     void endTensorLeaseScope(std::vector<c10::StorageImpl const*> const& outputStorages, int device,
         cudaStream_t stream, bool failed, std::optional<size_t> entryDepth = std::nullopt);
 
+    //! Release the leases of the calling thread's innermost scope on \p device on \p stream, as a
+    //! successful exit of that scope would, except those of \p keptStorages, which stay with the
+    //! scope. The scope stays open. Without an open scope, this releases nothing and logs a warning
+    //! the first time.
+    void releaseTensorLeaseScopeLeases(
+        std::vector<c10::StorageImpl const*> const& keptStorages, int device, cudaStream_t stream);
+
     // Destructor fallback. The lease ID prevents a late tensor destructor from releasing a newer
     // checkout of the same allocation. Fallback-released buffers are quarantined until communicator
     // teardown because destructor timing and stream identity are not rank-deterministic.
@@ -347,6 +357,12 @@ class NCCLWindowAllocator
     // Get the number of buffers allocated for a communicator
     size_t getBufferCount(ncclComm_t comm) const;
 
+    //! Whether this process has made a window request while NCCL windows are supported and allocation
+    //! is not suppressed on the calling thread, whatever the request's outcome. No tensor can hold a
+    //! window lease before the first such request. Every rank makes its requests at the same points
+    //! (see requestBuffer), so the answer changes at the same point on every rank. It never reverts.
+    bool hasRequestedBuffers() const noexcept;
+
     // Get the number of buffers in use for a communicator
     size_t getBufferInUseCount(ncclComm_t comm) const;
 
@@ -453,6 +469,7 @@ class NCCLWindowAllocator
     bool const mCaptureRegistrationSupported;
     mutable std::mutex mMutex;
     std::atomic_bool mSuppressLifecycleWarningAssertionsForTest{false};
+    std::atomic_bool mBufferRequested{false};
     std::unordered_map<ncclComm_t, std::vector<BufferEntry>> mBufferPool;
     std::unordered_set<ncclComm_t> mRegisteredComms;
     // Smallest request size that is known to fail collectively for each communicator.

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -1703,6 +1703,19 @@ void endNCCLWindowTensorScope(
     endTrackedNCCLWindowTensorScope(inputs, outputs, failed, -1);
 }
 
+void releaseNCCLWindowTensorScopeLeases(
+    torch::List<torch::Tensor> const& inputs, torch::List<torch::Tensor> const& outputs)
+{
+#if ENABLE_MULTI_DEVICE
+    auto const device = getTensorScopeInputs(inputs).second;
+    tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().releaseTensorLeaseScopeLeases(
+        getTensorStorages(outputs), device, at::cuda::getCurrentCUDAStream(device));
+#else
+    (void) inputs;
+    (void) outputs;
+#endif
+}
+
 namespace
 {
 
@@ -2537,10 +2550,13 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
     m.def("begin_nccl_window_tensor_scope.tracked(Tensor[] inputs) -> int");
     m.def(
         "end_nccl_window_tensor_scope.tracked(Tensor[] inputs, Tensor[] outputs, bool failed, int entry_depth) -> ()");
+    // Releases the innermost scope's leases, except the outputs', as its exit would; the scope stays open.
+    m.def("release_nccl_window_tensor_scope_leases(Tensor[] inputs, Tensor[] outputs) -> ()");
     m.def("set_nccl_window_graph_owner(int owner) -> ()");
     m.def("release_nccl_window_graph_owner(int owner) -> ()");
     m.def("set_nccl_window_allocation_suppressed(bool suppressed) -> bool");
     m.def("get_nccl_window_buffer_release_epoch() -> int");
+    m.def("has_requested_nccl_window_buffers() -> bool");
     m.def("promote_nccl_window_buffer_releases(int device, int release_epoch) -> ()");
     m.def(
         "minimax_allreduce_rms("
@@ -2581,6 +2597,7 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
     m.impl("end_nccl_window_tensor_scope", &tensorrt_llm::torch_ext::endNCCLWindowTensorScope);
     m.impl("begin_nccl_window_tensor_scope.tracked", &tensorrt_llm::torch_ext::beginTrackedNCCLWindowTensorScope);
     m.impl("end_nccl_window_tensor_scope.tracked", &tensorrt_llm::torch_ext::endTrackedNCCLWindowTensorScope);
+    m.impl("release_nccl_window_tensor_scope_leases", &tensorrt_llm::torch_ext::releaseNCCLWindowTensorScopeLeases);
     m.impl("is_nccl_window_buffer", &tensorrt_llm::torch_ext::isNCCLWindowBuffer);
     m.impl("minimax_allreduce_rms", &tensorrt_llm::torch_ext::minimax_allreduce_rms);
     m.impl("minimax_allreduce_rms_qk", &tensorrt_llm::torch_ext::minimax_allreduce_rms_qk);
@@ -2617,6 +2634,8 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
             return static_cast<int64_t>(
                 tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().getBufferReleaseEpoch());
         });
+    m.impl("has_requested_nccl_window_buffers",
+        []() { return tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().hasRequestedBuffers(); });
     m.impl("promote_nccl_window_buffer_releases",
         [](int64_t device, int64_t releaseEpoch)
         {
@@ -2629,6 +2648,7 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
     m.impl("release_nccl_window_graph_owner", [](int64_t) {});
     m.impl("set_nccl_window_allocation_suppressed", [](bool) { return false; });
     m.impl("get_nccl_window_buffer_release_epoch", []() { return int64_t{0}; });
+    m.impl("has_requested_nccl_window_buffers", []() { return false; });
     m.impl("promote_nccl_window_buffer_releases", [](int64_t, int64_t) {});
 #endif
 }
@@ -2652,4 +2672,6 @@ TORCH_LIBRARY_IMPL(trtllm, CPU, m)
         [](torch::List<torch::Tensor> const&, torch::List<torch::Tensor> const&, bool, int64_t) { return; });
     m.impl("end_nccl_window_tensor_scope",
         [](torch::List<torch::Tensor> const&, torch::List<torch::Tensor> const&, bool) { return; });
+    m.impl("release_nccl_window_tensor_scope_leases",
+        [](torch::List<torch::Tensor> const&, torch::List<torch::Tensor> const&) { return; });
 }

cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp

@@ -54,6 +54,11 @@ class NCCLWindowAllocatorTestAccess
         allocator.recordSymmetricFailureLocked(comm, size);
     }
 
+    static bool exchangeBufferRequested(NCCLWindowAllocator& allocator, bool requested)
+    {
+        return allocator.mBufferRequested.exchange(requested);
+    }
+
     static cudaError_t clearCudaErrorIfSymmetricAllocationFailed(
         int localAllocOk, NCCLWindowAllocator::CudaGetLastErrorFunc getLastError = cudaGetLastError)
     {
@@ -500,6 +505,44 @@ TEST_F(NCCLWindowAllocatorTest, BasicAllocation)
     allocator.releaseBuffer(*mComm, buffer.ptr);
 }
 
+TEST_F(NCCLWindowAllocatorTest, RequestedBuffersFollowUnsuppressedRequests)
+{
+    using Access = nccl_util::NCCLWindowAllocatorTestAccess;
+    auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();
+    auto const testComm = createWindowComm(*mComm, 0, mRank);
+    size_t constexpr kBufferSize = 4 * 1024;
+
+    // The flag is process-wide: start from an unrequested state and restore it on every exit.
+    struct RestoreRequested
+    {
+        nccl_util::NCCLWindowAllocator& allocator;
+        bool requested;
+
+        ~RestoreRequested()
+        {
+            Access::exchangeBufferRequested(allocator, requested);
+        }
+    } const restore{allocator, Access::exchangeBufferRequested(allocator, false)};
+
+    // Suppressed sections are not SPMD, so their requests are not recorded.
+    bool const previouslySuppressed = allocator.setWindowAllocationSuppressed(true);
+    EXPECT_FALSE(allocator.requestBuffer(*testComm, kBufferSize).isValid());
+    allocator.setWindowAllocationSuppressed(previouslySuppressed);
+    EXPECT_FALSE(allocator.hasRequestedBuffers());
+
+    auto const buffer = allocator.requestBuffer(*testComm, kBufferSize);
+    ASSERT_TRUE(buffer.isValid());
+    EXPECT_TRUE(allocator.hasRequestedBuffers());
+    allocator.releaseBuffer(*testComm, buffer.ptr);
+    EXPECT_TRUE(allocator.hasRequestedBuffers());
+
+    // A request that gets no buffer is recorded too: every rank makes it.
+    Access::exchangeBufferRequested(allocator, false);
+    Access::recordSymmetricFailure(allocator, *testComm, 2 * kBufferSize);
+    EXPECT_FALSE(allocator.requestBuffer(*testComm, 2 * kBufferSize).isValid());
+    EXPECT_TRUE(allocator.hasRequestedBuffers());
+}
+
 TEST_F(NCCLWindowAllocatorTest, ExplicitStreamAllocationAndCapture)
 {
     auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();
@@ -1387,6 +1430,61 @@ TEST_F(CreateNCCLWindowTensorTest, TensorLeaseScopeReleasesTemporariesAndEscapes
     EXPECT_TRUE(allocator.releaseTensorBuffer(output.storage().unsafeGetStorageImpl(), stream));
 }
 
+TEST_F(CreateNCCLWindowTensorTest, ScopeLeaseReleaseKeepsScopeAndKeptStorages)
+{
+    using nccl_util::createNCCLWindowTensor;
+
+    auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();
+    int const device = at::cuda::current_device();
+    cudaStream_t const stream = at::cuda::getCurrentCUDAStream();
+    std::vector<int64_t> shape = {16, 16};
+
+    allocator.beginTensorLeaseScope({}, device);
+    nccl_util::NCCLWindowBuffer temporaryBuffer;
+    {
+        auto [temporary, buffer] = createNCCLWindowTensor(mComm, shape, torch::kFloat32);
+        temporaryBuffer = buffer;
+        ASSERT_TRUE(temporaryBuffer.isValid());
+    }
+    auto [live, liveBuffer] = createNCCLWindowTensor(mComm, shape, torch::kFloat32);
+    auto [output, outputBuffer] = createNCCLWindowTensor(mComm, shape, torch::kFloat32);
+    ASSERT_TRUE(liveBuffer.isValid());
+    ASSERT_TRUE(outputBuffer.isValid());
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 3);
+
+    // As at a successful scope exit, the dead temporary and the live tensor that is not kept become
+    // reusable on the stream. The kept output stays with the scope, which stays open.
+    allocator.releaseTensorLeaseScopeLeases({output.storage().unsafeGetStorageImpl()}, device, stream);
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 1);
+    EXPECT_TRUE(
+        nccl_util::NCCLWindowAllocatorTestAccess::isReusableOnStream(allocator, *mComm, temporaryBuffer.ptr, stream));
+    EXPECT_TRUE(
+        nccl_util::NCCLWindowAllocatorTestAccess::isReusableOnStream(allocator, *mComm, liveBuffer.ptr, stream));
+
+    allocator.endTensorLeaseScope({}, device, stream, false);
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 0);
+    EXPECT_TRUE(
+        nccl_util::NCCLWindowAllocatorTestAccess::isReusableOnStream(allocator, *mComm, outputBuffer.ptr, stream));
+}
+
+TEST_F(CreateNCCLWindowTensorTest, ScopeLeaseReleaseWithoutScopeReleasesNothing)
+{
+    using nccl_util::createNCCLWindowTensor;
+
+    auto& allocator = nccl_util::NCCLWindowAllocator::getInstance();
+    int const device = at::cuda::current_device();
+    cudaStream_t const stream = at::cuda::getCurrentCUDAStream();
+    auto [tensor, buffer] = createNCCLWindowTensor(mComm, {16, 16}, torch::kFloat32);
+    ASSERT_TRUE(buffer.isValid());
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 1);
+
+    // No scope is open, so the unscoped lease stays in use.
+    allocator.releaseTensorLeaseScopeLeases({}, device, stream);
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 1);
+    releaseTensor(tensor);
+    EXPECT_EQ(allocator.getBufferInUseCount(*mComm), 0);
+}
+
 TEST_F(CreateNCCLWindowTensorTest, AuxiliaryStreamTensorJoinsCurrentThreadScope)
 {
     using nccl_util::createNCCLWindowTensor;

tensorrt_llm/_torch/nccl_window_tensor_scope.py

@@ -271,14 +271,47 @@ def remove(self) -> None:
             module.forward = self._previous
 
 
+def _release_first_request_leases(args: tuple, kwargs: dict[str, Any], output: Any) -> None:
+    """Release what an unscoped call that made the first window request leased, but its outputs.
+
+    No tensor holds a window lease before the process's first request, so every lease of the
+    innermost scope comes from that call. Releasing them is what a scope around the call would
+    have done at its exit.
+    """
+    inputs = _module_call_cuda_tensors(args, kwargs)
+    if not inputs:
+        # A call without CUDA inputs would not have entered a scope either.
+        return
+    outputs: list[torch.Tensor] = []
+    _output_cuda_tensors(output, outputs)
+    torch.ops.trtllm.release_nccl_window_tensor_scope_leases(inputs, outputs)
+
+
 def _install_forward_scope(module: torch.nn.Module) -> _ForwardScopeHandle:
     # A per-instance forward keeps nn.Module calls on their hook-free fast path; eager
     # decoder layers run once per layer per iteration.
     previous = module.__dict__.get("forward")
     forward = module.forward
+    window_buffers_requested = False
 
     @wraps(forward)
     def forward_in_scope(*args: Any, **kwargs: Any) -> Any:
+        nonlocal window_buffers_requested
+        # Before the process requests an NCCL window buffer, no input carries a window lease
+        # and a layer scope has nothing to adopt or release. Every rank makes these requests
+        # from the same operations, so all ranks start scoping at the same call and scope
+        # every call after it.
+        if not window_buffers_requested:
+            window_buffers_requested = torch.ops.trtllm.has_requested_nccl_window_buffers()
+            if not window_buffers_requested:
+                output = forward(*args, **kwargs)
+                window_buffers_requested = torch.ops.trtllm.has_requested_nccl_window_buffers()
+                if window_buffers_requested:
+                    # This call made the first request. Its window temporaries are released
+                    # now, as its layer scope would have released them, so the next layers
+                    # reuse their buffers.
+                    _release_first_request_leases(args, kwargs, output)
+                return output
         scope = _NCCLWindowTensorScope._for_module_call(args, kwargs)
         with scope:
             return scope.escape(forward(*args, **kwargs))
@@ -290,7 +323,12 @@ def forward_in_scope(*args: Any, **kwargs: Any) -> Any:
 def install_eager_nccl_window_tensor_scopes(
     model: torch.nn.Module,
 ) -> list[_ForwardScopeHandle]:
-    """Install eager model scopes without exposing Python hooks to Dynamo."""
+    """Install eager model scopes without exposing Python hooks to Dynamo.
+
+    Run the model inside ``nccl_window_tensor_scope``, as ``ModelCaller`` does. Layer calls run
+    unscoped until the process's first window request, and the call that makes it releases its
+    window temporaries from that enclosing scope.
+    """
     return [
         _install_forward_scope(module)
         for module in model.modules()