[None][fix] gate NCCL window autotune pre-hook

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -1541,6 +1541,47 @@ void preallocateNCCLWindowBuffer(
 #endif
 }
 
+bool isNCCLWindowBuffer(torch::Tensor const& input, torch::List<int64_t> const& group)
+{
+#if ENABLE_MULTI_DEVICE
+    if (!input.is_cuda() || input.numel() == 0 || group.size() == 0)
+    {
+        return false;
+    }
+
+    std::set<int> groupSet;
+    for (auto const& rank : group)
+    {
+        groupSet.insert(static_cast<int>(rank));
+    }
+
+    std::shared_ptr<ncclComm_t> commPtr;
+    try
+    {
+        commPtr = getComm(groupSet);
+    }
+    catch (std::exception const& e)
+    {
+        TLLM_LOG_DEBUG("[isNCCLWindowBuffer] getComm threw (MPI disabled?): %s", e.what());
+        return false;
+    }
+
+    if (!commPtr || *commPtr == nullptr)
+    {
+        TLLM_LOG_DEBUG("[isNCCLWindowBuffer] NCCL comm is null");
+        return false;
+    }
+
+    using tensorrt_llm::common::nccl_util::NCCLWindowAllocator;
+    auto const buffer = NCCLWindowAllocator::getInstance().searchBuffer(*commPtr, input.data_ptr());
+    return buffer.isValid();
+#else
+    (void) input;
+    (void) group;
+    return false;
+#endif
+}
+
 std::vector<torch::Tensor> allreduce_raw(torch::Tensor const& input, torch::optional<torch::Tensor> const& residual,
     torch::optional<torch::Tensor> const& norm_weight, torch::optional<torch::Tensor> const& scale,
     torch::optional<torch::Tensor> const& bias, torch::optional<torch::Tensor> workspace,
@@ -1982,6 +2023,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "int nranks,"
         "float eps) -> Tensor[]");
     m.def("preallocate_nccl_window_buffer(Tensor input, int[] group, int count) -> ()");
+    m.def("is_nccl_window_buffer(Tensor input, int[] group) -> bool");
     m.def(
         "minimax_allreduce_rms("
         "Tensor input,"
@@ -2012,6 +2054,7 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
     m.impl("moe_allreduce", &tensorrt_llm::torch_ext::moe_allreduce);
     m.impl("moe_finalize_allreduce", &tensorrt_llm::torch_ext::moe_finalize_allreduce);
     m.impl("preallocate_nccl_window_buffer", &tensorrt_llm::torch_ext::preallocateNCCLWindowBuffer);
+    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);
 }
@@ -2026,4 +2069,5 @@ TORCH_LIBRARY_IMPL(trtllm, CPU, m)
             return std::vector<at::Tensor>{};
         });
     m.impl("preallocate_nccl_window_buffer", [](at::Tensor const&, torch::List<int64_t> const&, int64_t) { return; });
+    m.impl("is_nccl_window_buffer", [](at::Tensor const&, torch::List<int64_t> const&) { return false; });
 }

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -1829,17 +1829,20 @@ def __init__(
         op: int,
         eps: float,
         trigger_completion_at_end: bool,
+        input_uses_nccl_window: bool = False,
     ):
         self.tp_size = tp_size
         self.op = op
         self.group = group
         self.eps = eps
         self.trigger_completion_at_end = trigger_completion_at_end
+        self.input_uses_nccl_window = input_uses_nccl_window
 
     def unique_id(self):
         return (
             self.tp_size,
             self.op,
+            self.input_uses_nccl_window,
         )
 
     @classmethod
@@ -1976,6 +1979,16 @@ def _(
     return None
 
 
+# Host-side predicate for custom-op implementations only. Do not call this
+# directly from model forward code or compiled graphs; it returns a Python bool
+# and is intended only to let another custom op gate its setup logic.
+@torch.library.register_fake("trtllm::is_nccl_window_buffer")
+def _(input: torch.Tensor, group: List[int]) -> bool:
+    raise AssertionError(
+        "trtllm::is_nccl_window_buffer should only be called inside another "
+        "custom op runtime implementation, not from fake/tracing execution.")
+
+
 @torch.library.custom_op("trtllm::tunable_allreduce", mutates_args=())
 def tunable_allreduce(
     input: torch.Tensor,
@@ -1992,35 +2005,51 @@ def tunable_allreduce(
 ) -> List[torch.Tensor]:
 
     tuner = AutoTuner.get()
+    group_list = list(group)
 
+    def _uses_nccl_symmetric_memory_window(input_tensor: torch.Tensor) -> bool:
+        # Keep is_nccl_window_buffer scoped inside this custom op. Calling it
+        # from normal model code would expose a host-side bool predicate to the
+        # graph instead of using it only to configure autotune execution.
+        return (isinstance(input_tensor, torch.Tensor) and input_tensor.is_cuda
+                and torch.ops.trtllm.is_nccl_window_buffer(
+                    input_tensor, group_list))
+
+    input_uses_nccl_window = _uses_nccl_symmetric_memory_window(input)
     allreduce_runner = AllReduceRunner(
         len(group),
         group,
         op,
         eps,
         trigger_completion_at_end,
+        input_uses_nccl_window,
     )
 
     def _inputs_pre_hook_register_nccl_symmetric_memory_window(
             inputs: List[torch.Tensor]) -> List[torch.Tensor]:
         if not inputs:
             return inputs
         input_tensor = inputs[0]
+        if not input_uses_nccl_window:
+            return inputs
         if not isinstance(input_tensor,
                           torch.Tensor) or not input_tensor.is_cuda:
             return inputs
         nccl_symmetric_memory_window_tensor, actual_kind = torch.ops.trtllm.allocate_output(
-            input_tensor, int(BufferKind.NCCL_WINDOW), list(group))
+            input_tensor, int(BufferKind.NCCL_WINDOW), group_list)
         if actual_kind != int(BufferKind.NCCL_WINDOW):
             return inputs
         nccl_symmetric_memory_window_tensor.copy_(input_tensor)
         new_inputs = list(inputs)
         new_inputs[0] = nccl_symmetric_memory_window_tensor
         return new_inputs
 
-    tuning_config = replace(
-        AllReduceRunner.tuning_config,
-        inputs_pre_hook=_inputs_pre_hook_register_nccl_symmetric_memory_window)
+    tuning_config = AllReduceRunner.tuning_config
+    if input_uses_nccl_window:
+        tuning_config = replace(
+            AllReduceRunner.tuning_config,
+            inputs_pre_hook=
+            _inputs_pre_hook_register_nccl_symmetric_memory_window)
 
     _, best_tactic = tuner.choose_one(
         "trtllm::tunable_allreduce::allreduce",