From 444085744cee737a164a7d2a0f79c192739cf6c2 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Tue, 26 May 2026 01:30:56 -0700 Subject: [PATCH] [None][fix] gate NCCL window autotune pre-hook The allreduce autotuner can profile with dummy tensors. When the real input is not an NCCL symmetric window buffer, registering a dummy window input during autotune lets NCCL_SYMMETRIC look valid even though runtime later sees a regular tensor and pays the extra copy or falls back. Add an NCCL-window buffer detector and include the real input window state in the allreduce runner identity. Only install the autotune pre-hook that allocates a symmetric window input when the real allreduce input already uses the NCCL symmetric memory window. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 44 +++++++++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 37 ++++++++++++++-- 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 3a6091022acd..e79c7fce26d5 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -1541,6 +1541,47 @@ void preallocateNCCLWindowBuffer( #endif } +bool isNCCLWindowBuffer(torch::Tensor const& input, torch::List const& group) +{ +#if ENABLE_MULTI_DEVICE + if (!input.is_cuda() || input.numel() == 0 || group.size() == 0) + { + return false; + } + + std::set groupSet; + for (auto const& rank : group) + { + groupSet.insert(static_cast(rank)); + } + + std::shared_ptr 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 allreduce_raw(torch::Tensor const& input, torch::optional const& residual, torch::optional const& norm_weight, torch::optional const& scale, torch::optional const& bias, torch::optional 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{}; }); m.impl("preallocate_nccl_window_buffer", [](at::Tensor const&, torch::List const&, int64_t) { return; }); + m.impl("is_nccl_window_buffer", [](at::Tensor const&, torch::List const&) { return false; }); } diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 7174484bc103..8a23b7958a32 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/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,13 +2005,24 @@ 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( @@ -2006,11 +2030,13 @@ def _inputs_pre_hook_register_nccl_symmetric_memory_window( 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) @@ -2018,9 +2044,12 @@ def _inputs_pre_hook_register_nccl_symmetric_memory_window( 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",