[https://nvbugs/6944374][fix] Fix rank-asymmetric NCCL window allocation hang during PARALLEL autotuning

cpp/tensorrt_llm/common/ncclUtils.cpp

@@ -24,6 +24,7 @@
 #include <cstdlib>
 #include <limits>
 #include <stdexcept>
+#include <utility>
 
 namespace
 {
@@ -398,6 +399,7 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept
 namespace
 {
 thread_local std::optional<uint64_t> gGraphPoolOwner;
+thread_local bool gWindowAllocationSuppressed = false;
 } // namespace
 
 NCCLWindowAllocator& NCCLWindowAllocator::getInstance()
@@ -417,6 +419,11 @@ void NCCLWindowAllocator::setGraphPoolOwner(int64_t owner)
     gGraphPoolOwner = owner < 0 ? std::nullopt : std::optional<uint64_t>{static_cast<uint64_t>(owner)};
 }
 
+bool NCCLWindowAllocator::setWindowAllocationSuppressed(bool suppressed)
+{
+    return std::exchange(gWindowAllocationSuppressed, suppressed);
+}
+
 void NCCLWindowAllocator::releaseGraphPoolOwner(int64_t owner)
 {
     TLLM_CHECK_WITH_INFO(owner >= 0, "NCCL window graph-pool owner must be non-negative");
@@ -469,6 +476,12 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size
     {
         return NCCLWindowBuffer();
     }
+    if (gWindowAllocationSuppressed)
+    {
+        TLLM_LOG_TRACE("[NCCLUtil] NCCL window allocation suppressed on this thread for comm %p (requested: %zu)",
+            static_cast<void*>(comm), size);
+        return NCCLWindowBuffer();
+    }
 
     TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null");
     TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0");

cpp/tensorrt_llm/common/ncclUtils.h

@@ -245,6 +245,11 @@ class NCCLWindowAllocator
     // still in use are returned when their deterministic tensor scope ends.
     void releaseGraphPoolOwner(int64_t owner);
 
+    // While set on this host thread, requestBuffer returns an invalid buffer without touching the
+    // pool, so callers fall back to plain memory. Used where ranks are not SPMD, such as PARALLEL
+    // autotuning. Returns the previous setting.
+    bool setWindowAllocationSuppressed(bool suppressed);
+
     // Search for a buffer by pointer. Returns an invalid buffer if not found.
     // This matches the UBManager.search_buffer() interface.
     NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const;

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -2498,6 +2498,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "end_nccl_window_tensor_scope.tracked(Tensor[] inputs, Tensor[] outputs, bool failed, int entry_depth) -> ()");
     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("promote_nccl_window_buffer_releases(int device, int release_epoch) -> ()");
     m.def(
@@ -2558,6 +2559,12 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
     m.impl("release_nccl_window_graph_owner",
         [](int64_t owner)
         { tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().releaseGraphPoolOwner(owner); });
+    m.impl("set_nccl_window_allocation_suppressed",
+        [](bool suppressed)
+        {
+            return tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().setWindowAllocationSuppressed(
+                suppressed);
+        });
     m.impl("get_nccl_window_buffer_release_epoch",
         []()
         {
@@ -2574,6 +2581,7 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
 #else
     m.impl("set_nccl_window_graph_owner", [](int64_t) {});
     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("promote_nccl_window_buffer_releases", [](int64_t, int64_t) {});
 #endif

tensorrt_llm/_torch/autotuner.py

@@ -34,8 +34,8 @@
 import tensorrt_llm
 from tensorrt_llm._torch.distributed import Distributed
 from tensorrt_llm._torch.nccl_window_graph import (
-    abandon_nccl_window_graph_owner, nccl_window_graph_capture,
-    release_nccl_window_graph_owner)
+    abandon_nccl_window_graph_owner, nccl_window_allocation_suppressed,
+    nccl_window_graph_capture, release_nccl_window_graph_owner)
 from tensorrt_llm._torch.nccl_window_tensor_scope import \
     nccl_window_tensor_scope
 from tensorrt_llm._utils import confidential_compute_enabled, nvtx_range
@@ -1332,7 +1332,14 @@ def _profile_runners(
         tuning_config: TuningConfig,
         **kwargs,
     ) -> float:
-        with nccl_window_tensor_scope(input_tensors) as scope:
+        # PARALLEL ranks profile different tactics, so window requests here would diverge the
+        # per-rank pools and let a later request register a window collectively on some ranks only.
+        window_guard = (nccl_window_allocation_suppressed()
+                        if tuning_config.distributed_tuning_strategy
+                        == DistributedTuningStrategy.PARALLEL
+                        and self._is_distributed() else
+                        contextlib.nullcontext())
+        with window_guard, nccl_window_tensor_scope(input_tensors) as scope:
             scope.escape(input_tensors)
             if tuning_config.inputs_pre_hook is not None:
                 input_tensors = tuning_config.inputs_pre_hook(input_tensors)

tensorrt_llm/_torch/nccl_window_graph.py

@@ -106,3 +106,18 @@ def release_nccl_window_graph_owner(pool: Any) -> None:
         if owner is not None:
             torch.ops.trtllm.release_nccl_window_graph_owner(owner)
             del _pool_owners[key]
+
+
+@contextlib.contextmanager
+def nccl_window_allocation_suppressed() -> Iterator[None]:
+    """Serve NCCL-window requests on this thread from plain memory.
+
+    The window pool is left untouched, so work that does not run in SPMD order
+    across ranks cannot make a later request miss on only some ranks and enter
+    the collective registration alone.
+    """
+    previous = torch.ops.trtllm.set_nccl_window_allocation_suppressed(True)
+    try:
+        yield
+    finally:
+        torch.ops.trtllm.set_nccl_window_allocation_suppressed(previous)

tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py

@@ -29,6 +29,13 @@
 
 import tensorrt_llm
 import tensorrt_llm.bindings.internal.userbuffers as ub
+from tensorrt_llm._torch.autotuner import (
+    AutoTuner,
+    DistributedTuningStrategy,
+    TunableRunner,
+    TuningConfig,
+    autotune,
+)
 from tensorrt_llm._torch.distributed import (
     AllReduce,
     AllReduceFusionOp,
@@ -222,6 +229,56 @@ def _allocate(ref, kind, tp_group):
     return True
 
 
+def _run_parallel_tuning_then_window_allreduce(tp_size, tp_rank, seq_len, hidden_size):
+    """PARALLEL-tune a window-output runner, then all-reduce in a fresh window on every rank."""
+    from tensorrt_llm.bindings.internal.thop import BufferKind
+
+    def _allocate(ref, tp_group):
+        return torch.ops.trtllm.allocate_output(ref, int(BufferKind.NCCL_WINDOW), tp_group)
+
+    class WindowOutputRunner(TunableRunner):
+        def __init__(self, tp_group):
+            self.tp_group = tp_group
+
+        def get_valid_tactics(self, inputs, profile, **kwargs):
+            return [0]
+
+        def forward(self, inputs, *, tactic=-1, **kwargs):
+            out, _ = _allocate(inputs[0], self.tp_group)
+            return out.copy_(inputs[0])
+
+    mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=tp_rank)
+    ref = torch.ones(seq_len, hidden_size, dtype=torch.bfloat16, device="cuda")
+    with discard_nccl_window_tensor_outputs(ref):
+        _, actual_kind_int = _allocate(ref, mapping.tp_group)
+    if actual_kind_int != int(BufferKind.NCCL_WINDOW):
+        return _SKIP
+
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    tuner.setup_distributed_state(mapping)
+    try:
+        # A single tactic is profiled by rank 0 only.
+        with autotune():
+            tuner.choose_one(
+                "test_parallel_tuning_then_window_allreduce",
+                [WindowOutputRunner(mapping.tp_group)],
+                TuningConfig(distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL),
+                [ref],
+            )
+    finally:
+        tuner.setup_distributed_state(Mapping(), None)
+
+    allreduce = AllReduce(mapping=mapping, strategy=AllReduceStrategy.NCCL_SYMMETRIC)
+    with discard_nccl_window_tensor_outputs(ref):
+        out, actual_kind_int = _allocate(ref, mapping.tp_group)
+        assert actual_kind_int == int(BufferKind.NCCL_WINDOW)
+        out.copy_(ref)
+        result = allreduce(out, all_reduce_params=AllReduceParams(fusion_op=AllReduceFusionOp.NONE))
+        torch.testing.assert_close(result, torch.full_like(result, tp_size), atol=1e-2, rtol=1e-2)
+    return True
+
+
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
 @pytest.mark.parametrize(
     "seq_len,hidden_size",
@@ -297,3 +354,18 @@ def test_allreduce_nccl_window_buffer(seq_len, hidden_size, mpi_pool_executor):
     )
     for r in results:
         assert r is True or r == _SKIP, f"Unexpected result: {r}"
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+@pytest.mark.timeout(300)
+@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
+def test_parallel_tuning_keeps_nccl_windows_symmetric(mpi_pool_executor):
+    tp_size = mpi_pool_executor.num_workers
+    results = list(
+        mpi_pool_executor.map(
+            run_single_rank,
+            *zip(*[(tp_size, _run_parallel_tuning_then_window_allreduce, 256, 4096)] * tp_size),
+        )
+    )
+    for r in results:
+        assert r is True or r == _SKIP, f"Unexpected result: {r}"