From dfcec555403b4c420de264d90936347836d57698 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 22 Jan 2026 12:16:40 +0000 Subject: [PATCH 1/2] [https://nvbugs/5814309][fix] Use NCCL as fallback to avoid crash due to insufficient memory. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- cpp/tensorrt_llm/common/customAllReduceUtils.h | 2 +- cpp/tensorrt_llm/thop/allreduceOp.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/common/customAllReduceUtils.h b/cpp/tensorrt_llm/common/customAllReduceUtils.h index 5f35ff5f89e8..f69eb0cb2565 100644 --- a/cpp/tensorrt_llm/common/customAllReduceUtils.h +++ b/cpp/tensorrt_llm/common/customAllReduceUtils.h @@ -115,7 +115,7 @@ inline AllReduceStrategyType selectStrategyLookUpTable( || num_token_index >= AllReduceBestStrategyTable.at(sm_version).at(tp_index).at(fusion_op_index).at(hidden_size_index).size()) { - return AllReduceStrategyType::NCCL_SYMMETRIC; + return AllReduceStrategyType::NCCL; } return static_cast( diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index c75324251865..a35fdfbc2e5b 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -1206,7 +1206,7 @@ class AllreduceOp if (ifFallbackToNCCL(seq_len, message_size_bytes, max_workspace_size)) { - return AllReduceStrategyType::NCCL_SYMMETRIC; + return AllReduceStrategyType::NCCL; } // This rule based heuristic only chooses between NCCL_SYMMETRIC and MIN_LATENCY strategies. From 91c49b7ac12c8179bb9e6bc3110cec301bf55bd7 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 22 Jan 2026 12:33:29 +0000 Subject: [PATCH 2/2] remove unnecessary noexcept Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index a35fdfbc2e5b..903908288756 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -281,7 +281,7 @@ class AllreduceOp std::vector run(torch::Tensor const& input, torch::optional const& residual, torch::optional const& norm_weight, torch::optional const& scale, torch::optional const& bias, bool trigger_completion_at_end, - torch::optional workspace) noexcept + torch::optional workspace) { size_t size = input.numel(); size_t seq_len = input.size(0); @@ -564,7 +564,7 @@ class AllreduceOp std::vector runLowPrecisionAllReduce(torch::Tensor const& input, torch::optional const& residual, torch::optional const& norm_weight, - torch::optional const& scale, torch::optional const& bias) noexcept + torch::optional const& scale, torch::optional const& bias) { #ifdef ENABLE_FP8 auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); @@ -632,8 +632,7 @@ class AllreduceOp std::vector runFusionAllReduce(torch::Tensor const& input, torch::optional const& residual, torch::optional const& norm_weight, torch::optional const& scale, torch::optional const& bias, - bool trigger_completion_at_end, torch::optional workspace, - AllReduceStrategyType strategy) noexcept + bool trigger_completion_at_end, torch::optional workspace, AllReduceStrategyType strategy) { // Should handle only Lamport implementation auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); @@ -1232,7 +1231,8 @@ class AllreduceOp bool ifFallbackToNCCL(size_t seq_len, size_t message_size_bytes, size_t max_workspace_size) { - // If messageSize is greater than maxWorkspaceSize or topology is unsuitable, use NCCL_SYMMETRIC fallback. + // If messageSize is greater than maxWorkspaceSize or topology is unsuitable, use NCCL fallback. + // TODO: Use NCCL_SYMMETRIC once the memory allocation issue is resolved. if (message_size_bytes > max_workspace_size || !mIsP2PSupported || !mIsNVLINKSupported) { return true;