-
Notifications
You must be signed in to change notification settings - Fork 1.9k
[None][feat] Enable NCCL_SYMMETRIC as default fallback for AllReduce #9314
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
[None][feat] Enable NCCL_SYMMETRIC as default fallback for AllReduce #9314
Conversation
📝 WalkthroughWalkthroughThis pull request refactors NCCL-related infrastructure across TensorRT-LLM, introducing centralized NCCL utilities for dynamic library loading, resource management, and window buffer pooling. The changes migrate all-reduce strategies from NCCL to NCCL_SYMMETRIC as the default symmetric variant, remove legacy NCCL allocators, and add PyTorch integration for NCCL window tensors with comprehensive test coverage. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant Alloc as NCCLWindowAllocator
participant Mgr as NcclCommResourceManager
participant NCCL as NCCL Library
participant CUDA as CUDA
App->>Alloc: requestBuffer(comm, size)
Alloc->>Alloc: searchBuffer() - reuse if available
alt Buffer found
Alloc->>App: return NCCLWindowBuffer
else No buffer
Alloc->>CUDA: allocate device memory
CUDA-->>Alloc: device ptr
Alloc->>NCCL: ncclCommWindowRegister(ptr, size)
NCCL-->>Alloc: ncclWindow_t
Alloc->>Mgr: registerResource(comm, cleanup_callback)
Mgr->>Mgr: store cleanup in per-comm list
Alloc->>App: return NCCLWindowBuffer
end
App->>App: use buffer for allreduce
App->>Alloc: releaseBuffer(comm, ptr)
Alloc->>Alloc: mark buffer inUse=false
Note over Mgr,NCCL: On comm destruction
Mgr->>Mgr: cleanupResources(comm)
loop for each registered cleanup
Mgr->>Alloc: invoke cleanup callbacks
Alloc->>Alloc: cleanupBuffersForComm(comm)
Alloc->>NCCL: deregister windows
Alloc->>CUDA: free device memory
end
sequenceDiagram
participant Python as Python API
participant Torch as PyTorch
participant WindowTensor as ncclWindowTensor
participant Alloc as NCCLWindowAllocator
participant Comm as NCCL Comm
Python->>WindowTensor: create_nccl_window_tensor(group, shape, dtype)
WindowTensor->>Comm: getComm(group)
Comm-->>WindowTensor: ncclComm_t
WindowTensor->>Alloc: requestBuffer(comm, buffer_size)
Alloc-->>WindowTensor: NCCLWindowBuffer
WindowTensor->>Torch: create tensor from buffer ptr
Torch-->>WindowTensor: torch::Tensor with custom_deleter
Note over Torch: custom_deleter calls releaseBuffer on destruction
WindowTensor->>Python: return tensor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
This makes NCCL_SYMMETRIC default with some infra changes. The one thing, this does not yet include is the optimization that results are automatically written into NCCLWindowTensors, like the UB tensors. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/functional.py (1)
3981-4040: Review comment is valid; fix prevents passingNonetensor to TensorRT pluginThe review comment accurately identifies a logic flow issue. When
current_all_reduce_helper().workspace is Noneat line 4096, the code setsstrategy = AllReduceStrategy.NCCL_SYMMETRIC(line 4097) but leavesworkspace = None(line 4094). SinceNCCL_SYMMETRICis a distinct enum value (notNCCLorUB), the condition at line 4025 increate_allreduce_plugin()evaluates to true, appending theNoneworkspace toplug_inputs. ThisNonewould then be passed tonetwork.add_plugin_v2(), which expectsITensorinstances.The suggested fix—checking
if workspace is not Nonebefore appending at line 4025–4026—is the correct approach and prevents this failure path.cpp/tensorrt_llm/thop/allreduceOp.cpp (1)
437-537: Threshold overflow issue confirmed for ≥32-rank clustersThe review comment is accurate. Verification confirms:
Code location and coefficients verified: Exact match at line 478 of
cpp/tensorrt_llm/thop/allreduceOp.cppwith coefficients a = −4986.43478503, b = 156716.52177552Mathematical issue is real: The threshold becomes negative at ~31.43 ranks. For nRanks ≥ 32, the calculation yields negative values (e.g., −2849.39 at 32 ranks, −162415.30 at 64 ranks)
No guards exist: Search for threshold validation found zero safeguards—no clamping, max(), or bounds checking
Impact: Casting a negative double to
size_twraps to an extremely large positive value. This causes the comparison at line 495 to always evaluate true, silently disabling buffer registration for all buffer sizes on clusters with ≥32 ranks and severely degrading performance in a non-obvious, rank-dependent mannerThe suggested fix appropriately clamps the threshold to 0.0 before casting, preserving the tuning curve where it's valid and preventing the underflow.
🧹 Nitpick comments (12)
cpp/tensorrt_llm/common/customAllReduceUtils.h (1)
63-85: Unreachable NCCL_SYMMETRIC fallback inSelectStrategyLP
return AllReduceStrategyType::NCCL_SYMMETRIC;at Line 84 is currently unreachable because both branches above return. If this is meant as a real fallback for future conditions (e.g., an explicit “NCCL zone”), consider either:
- Adding an explicit
else/guard that can actually reach this return, or- Dropping the line (or adding a brief comment) to avoid confusion about dead code.
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2746-2764: NCCL_SYMMETRIC UB gating is correct; consider simplifying caller logicThe changes in
_init_userbufferscorrectly ensure that:
- TP size ≤ 1 and unsupported platforms still early-exit, and
self.llm_args.allreduce_strategy == "NCCL_SYMMETRIC"returnsFalsebefore callingub.initialize_userbuffers_manager, so NCCL_SYMMETRIC no longer sets up UB and can rely solely onNCCLWindowAllocator.Two small cleanups to consider:
In
__init__,use_ub_for_nccl = (self.llm_args.allreduce_strategy == "NCCL_SYMMETRIC" and self._init_userbuffers(...))will now always beFalse, since_init_userbuffersreturnsFalsefor NCCL_SYMMETRIC. This makesuse_ub_for_nccleffectively dead logic and causes_init_userbuffersto be called twice for NCCL_SYMMETRIC when torch.compile UB is enabled. You could:
- Skip calling
_init_userbuffersentirely whenallreduce_strategy == "NCCL_SYMMETRIC"in the caller, and/or- Rename or remove
use_ub_for_ncclto better reflect the new semantics.
use_nccl_symmetricis now hard-coded toFalsefor UB initialization. If there is no remaining UB path that depends on this flag, consider dropping the parameter fromub.initialize_userbuffers_manager(or at least the local variable) to avoid suggesting configurable behavior that no longer exists.These are readability/maintainability nits; behavior for NCCL_SYMMETRIC and other strategies looks correct.
tests/microbenchmarks/all_reduce.py (1)
171-183: AddingNCCL_SYMMETRICto benchmark strategies looks consistentIncluding
AllReduceStrategy.NCCL_SYMMETRICin the benchmark grid aligns with the new fallback behavior and will help compare it fairly against the existing modes. You may optionally consider skipping it when NCCL symmetric support is not available (e.g., via a helper similar to other NCCL feature checks), but it’s fine to leave responsibility to the runtime if that’s the established pattern.cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp (1)
17-49: Deprecation handling foruse_nccl_symmetricis clear and safeLogging a warning and otherwise ignoring
use_nccl_symmetrickeeps the API backward-compatible while reflecting the new NCCL_SYMMETRIC implementation that bypasses the userbuffer allocator. Longer term, you might consider removing this parameter from the public API once callers are migrated, but the current behavior is fine.cpp/tensorrt_llm/thop/ncclWindowTensor.h (1)
18-21: Drop unnecessaryncclUtils.hinclude from the header
ncclWindowTensor.honly needs Torch types for the function signature; it doesn’t use any symbols fromtensorrt_llm/common/ncclUtils.h. Including it here adds coupling and can trigger tooling issues (as seen in the static analysis hint).You can rely on the
.cppto include the NCCL utilities instead:-#include "tensorrt_llm/common/ncclUtils.h" -#include <torch/extension.h> -#include <vector> +#include <torch/extension.h> +#include <vector>This keeps the public declaration lightweight and avoids unnecessary rebuilds when NCCL internals change.
cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp (1)
336-341: Confirm behavior whenmStrategy == NCCL_SYMMETRICin plugin usage
enqueue()treatsruntimeStrategy == NCCLandNCCL_SYMMETRICequivalently for execution, which is fine. However:
supportsFormatCombination()only treatsNCCLandUBas single‑input (base_inputs = 1);NCCL_SYMMETRICis grouped with the “other strategies” path (base_inputs = 2).- In the fused NCCL path,
fusion_ptr_idxnow treatsmStrategy == NCCL_SYMMETRICas the single‑input layout (fusion_ptr_idx = 1), assuming NCCL‑style inputs.If a network ever constructs this plugin with
strategy == NCCL_SYMMETRIC(not just AUTO falling back at runtime), the input layout assumptions betweensupportsFormatCombination()andenqueue()may diverge.Either:
- Treat
NCCL_SYMMETRIClikeNCCLeverywhere layout‑wise (includingsupportsFormatCombination’sbase_inputs), or- Explicitly ensure plugin construction never passes
NCCL_SYMMETRICasmStrategyand keep it as a runtime‑only selection.Please double‑check intended usage and adjust one side for consistency.
Also applies to: 360-365, 383-391
cpp/tensorrt_llm/thop/allreduceOp.cpp (2)
95-149: CheckgetLocalGroup’s use ofLOCAL_COMM_SESSIONin the manual branchIn the
elsebranch (whengroup.size() < localSize), this implementation usesLOCAL_COMM_SESSION.send/recvwith*group.begin()and other world‑rank values as the destination/source ranks. In the plugin implementation, the equivalent manual branch usesCOMM_SESSIONfor these operations, which matches the fact thatgroupis in world‑rank space (seecpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp:getLocalGroup). Based on that precedent, usingLOCAL_COMM_SESSIONhere may be incorrect on multi‑node setups whengroup.size() < localSize.Please double‑check that:
LOCAL_COMM_SESSIONis defined such that using world ranks as peers is valid in this code path, or- If not, consider switching these manual send/recv calls back to
COMM_SESSIONto mirror the plugin’s behavior and avoid rank mismatches.
986-1040: Align comments and cleanup inselectImplementation/ifFallbackToNCCLA couple of small inconsistencies here:
ifFallbackToNCCL’s comment says “If messageSize is less than maxWorkspaceSize, use NCCL_SYMMETRIC…”, but the condition ismessage_size_bytes > max_workspace_size || !mIsP2PSupported || !mIsNVLINKSupported, i.e., fallback when the message is larger than the workspace or topology is unsuitable. The comment should reflect the actual predicate.- The final
return AllReduceStrategyType::NCCL_SYMMETRIC;at the end ofselectImplementationis unreachable because all preceding branches already return.I’d suggest updating the comment to match the logic and removing the dead return for clarity. Optionally, renaming
ifFallbackToNCCLto something likeshouldFallbackToNCCLBasedStrategyorshouldUseNcclSymmetricFallbackwould better reflect the new behavior but isn’t strictly necessary.tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py (3)
46-57: Simplify dynamic access tocreate_nccl_window_tensorInside
_create_nccl_window_tensor, you can drop the nestedgetattrcalls with constant names:func = getattr(getattr(_torch, "ops"), "trtllm").create_nccl_window_tensorand just write:
func = _torch.ops.trtllm.create_nccl_window_tensorYou still avoid storing a module‑level reference to
torch.ops(the function is resolved at call time), but the code is clearer and avoids the Ruff B009 warning.
99-200: Unusedtensor_parallel_rankparameters in helper testsSeveral helpers (
run_window_tensor_creation_test,run_window_tensor_multiple_test,run_window_tensor_different_shapes_test,run_window_tensor_operations_test) accepttensor_parallel_rankbut don’t use it, since only the AllReduce test needs the rank forMapping.Given these functions are invoked through a common wrapper signature, the extra parameter is understandable. If you’d like to quiet Ruff’s
ARG001warnings without changing call sites, you can rename the parameter to_tensor_parallel_rankor add a trivial use like_ = tensor_parallel_rankwith a comment indicating it is kept for signature consistency.Also applies to: 186-200, 221-236
363-387: Optional: considerzip(strict=True)if Python version allowsThe MPIPoolExecutor tests build argument lists via patterns like:
results = mpi_pool_executor.map( run_single_rank_test, *zip( *[ ( tensor_parallel_size, run_window_tensor_creation_test, shape, dtype_str, tensor_parallel_size, None, ) ] * tensor_parallel_size ), )Because all iterables arise from repeating the same tuple
tensor_parallel_sizetimes, their lengths are guaranteed equal. To satisfy Ruff’s B905 and make mismatches explicit if these patterns evolve, you could addstrict=Truewhen your minimum supported Python version includes it (3.10+):*zip( *[ ... ] * tensor_parallel_size, strict=True, )If Python < 3.10 must remain supported, the current code is logically correct and can be left as is.
Also applies to: 393-420, 425-452, 455-482, 485-512
cpp/tensorrt_llm/common/ncclUtils.h (1)
197-228:NCCLWindowBuffervalidity semantics differ slightly fromUBBuffer– confirm intended behaviorCompared to
runtime::ub::UBBuffer,NCCLWindowBuffer::isValid()additionally requireswindow != nullptr(Line 215). This is stricter thanUBBuffer::invalid()(which ignores the window field) and will cause buffers with a null window to be treated as invalid even ifptr/handle/sizeare set.If that’s intentional (i.e., a buffer is only usable once it’s fully window‑registered), this is fine and the design is clear. If you expect to stage allocations before registration, you may want a separate predicate (e.g.,
isAllocated()vsisRegistered()) so callers can distinguish between “no memory” and “registration missing.”
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (20)
cpp/tensorrt_llm/common/customAllReduceUtils.h(2 hunks)cpp/tensorrt_llm/common/ncclUtils.cpp(1 hunks)cpp/tensorrt_llm/common/ncclUtils.h(1 hunks)cpp/tensorrt_llm/common/opUtils.cpp(2 hunks)cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cpp(1 hunks)cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.h(0 hunks)cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp(2 hunks)cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp(5 hunks)cpp/tensorrt_llm/thop/CMakeLists.txt(1 hunks)cpp/tensorrt_llm/thop/allreduceOp.cpp(4 hunks)cpp/tensorrt_llm/thop/ncclWindowTensor.cpp(1 hunks)cpp/tensorrt_llm/thop/ncclWindowTensor.h(1 hunks)cpp/tests/unit_tests/multi_gpu/CMakeLists.txt(1 hunks)cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp(1 hunks)tensorrt_llm/_torch/pyexecutor/model_engine.py(1 hunks)tensorrt_llm/functional.py(1 hunks)tests/microbenchmarks/all_reduce.py(2 hunks)tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py(2 hunks)tests/unittest/_torch/multi_gpu/test_allreduce.py(1 hunks)tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py(1 hunks)
💤 Files with no reviewable changes (1)
- cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.h
🧰 Additional context used
🧠 Learnings (25)
📓 Common learnings
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device allreduce implementation (cpp/tensorrt_llm/thop/allreduceOp.cpp), the goto pattern in runNCCLAllReduceDeviceFusion is intentionally used for future extensibility, allowing multiple switch cases to fallback to the default handler. While not aesthetically ideal, this pattern supports adding more fusion cases later that can reuse the same fallback logic.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: tests/unittest/_torch/multi_gpu/test_nccl_device.py:138-149
Timestamp: 2025-10-13T19:45:03.518Z
Learning: In test_nccl_device.py, the NCCL device AllReduce implementation compares the entire residual tensor on each rank, unlike the UB implementation which compares per-rank chunks. The residual chunking calculations in the test are intentionally overridden to reflect this design difference.
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7763
File: cpp/tensorrt_llm/CMakeLists.txt:297-301
Timestamp: 2025-09-16T09:30:09.716Z
Learning: In the TensorRT-LLM project, NCCL libraries are loaded earlier by PyTorch libraries or the bindings library, so the main shared library doesn't need NCCL paths in its RPATH - the libraries will already be available in the process address space when needed.
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7520
File: tensorrt_llm/_torch/pyexecutor/resource_manager.py:605-613
Timestamp: 2025-09-24T03:31:28.908Z
Learning: In TensorRT-LLM Ray orchestrator mode, ProcessGroups are initialized with both Gloo and NCCL backends (e.g., "cuda:nccl,cpu:gloo"), allowing PyTorch distributed to automatically route CPU tensors through Gloo and GPU tensors through NCCL. This eliminates the need for manual device placement when performing allreduce operations on base types.
Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels, the <sstream> header is not needed as an explicit include in config.cu because it's provided transitively through other headers. Local compilation testing confirms this works without the explicit include.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/config.cu), std::ostringstream is used but <sstream> doesn't need to be explicitly included because it's provided transitively through other headers like tensorrt_llm/common/cudaUtils.h or config.h. Local compilation testing confirms this works without the explicit include.
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:399-417
Timestamp: 2025-08-21T21:48:35.135Z
Learning: CUTLASS extensions in TensorRT-LLM (located under cpp/tensorrt_llm/cutlass_extensions/) are designed to integrate with and extend functionality in the external CUTLASS repository. When analyzing these extensions, their consumers and functionality wiring may exist in the CUTLASS codebase rather than within TensorRT-LLM itself.
Learnt from: achartier
Repo: NVIDIA/TensorRT-LLM PR: 6763
File: tests/integration/defs/triton_server/conftest.py:16-22
Timestamp: 2025-08-11T20:09:24.389Z
Learning: In the TensorRT-LLM test infrastructure, the team prefers simple, direct solutions (like hard-coding directory traversal counts) over more complex but robust approaches when dealing with stable directory structures. They accept the maintenance cost of updating tests if the layout changes.
📚 Learning: 2025-09-23T15:01:00.070Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels, the <sstream> header is not needed as an explicit include in config.cu because it's provided transitively through other headers. Local compilation testing confirms this works without the explicit include.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tests/unit_tests/multi_gpu/CMakeLists.txtcpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.hcpp/tensorrt_llm/common/customAllReduceUtils.hcpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpptensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-23T15:01:00.070Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/config.cu), std::ostringstream is used but <sstream> doesn't need to be explicitly included because it's provided transitively through other headers like tensorrt_llm/common/cudaUtils.h or config.h. Local compilation testing confirms this works without the explicit include.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/thop/ncclWindowTensor.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.hcpp/tensorrt_llm/common/customAllReduceUtils.hcpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tests/unit_tests/multi_gpu/CMakeLists.txttests/unittest/_torch/multi_gpu/test_allreduce.pycpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/thop/ncclWindowTensor.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.hcpp/tensorrt_llm/common/customAllReduceUtils.htests/microbenchmarks/all_reduce.pycpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpptensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-16T09:30:09.716Z
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7763
File: cpp/tensorrt_llm/CMakeLists.txt:297-301
Timestamp: 2025-09-16T09:30:09.716Z
Learning: In the TensorRT-LLM project, NCCL libraries are loaded earlier by PyTorch libraries or the bindings library, so the main shared library doesn't need NCCL paths in its RPATH - the libraries will already be available in the process address space when needed.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tests/unit_tests/multi_gpu/CMakeLists.txtcpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/thop/ncclWindowTensor.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.hcpp/tensorrt_llm/common/customAllReduceUtils.htensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-10-13T19:45:03.518Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: tests/unittest/_torch/multi_gpu/test_nccl_device.py:138-149
Timestamp: 2025-10-13T19:45:03.518Z
Learning: In test_nccl_device.py, the NCCL device AllReduce implementation compares the entire residual tensor on each rank, unlike the UB implementation which compares per-rank chunks. The residual chunking calculations in the test are intentionally overridden to reflect this design difference.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/multi_gpu/test_allreduce.pycpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/common/customAllReduceUtils.htests/microbenchmarks/all_reduce.pytests/scripts/allreduce_perf/allreduce_heuristic_code_gen.pycpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpptensorrt_llm/functional.pytests/unittest/_torch/multi_gpu/test_ncclwindowtensor.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-23T14:58:05.372Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/thop/ncclWindowTensor.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.hcpp/tensorrt_llm/common/customAllReduceUtils.hcpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpptensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device allreduce implementation (cpp/tensorrt_llm/thop/allreduceOp.cpp), the goto pattern in runNCCLAllReduceDeviceFusion is intentionally used for future extensibility, allowing multiple switch cases to fallback to the default handler. While not aesthetically ideal, this pattern supports adding more fusion cases later that can reuse the same fallback logic.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txttensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/multi_gpu/test_allreduce.pycpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/common/customAllReduceUtils.htests/microbenchmarks/all_reduce.pytests/scripts/allreduce_perf/allreduce_heuristic_code_gen.pytensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-08-14T06:36:40.701Z
Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/multi_gpu/test_allreduce.pycpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/customAllReduceUtils.htensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cpp
📚 Learning: 2025-09-24T03:31:28.908Z
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7520
File: tensorrt_llm/_torch/pyexecutor/resource_manager.py:605-613
Timestamp: 2025-09-24T03:31:28.908Z
Learning: In TensorRT-LLM Ray orchestrator mode, ProcessGroups are initialized with both Gloo and NCCL backends (e.g., "cuda:nccl,cpu:gloo"), allowing PyTorch distributed to automatically route CPU tensors through Gloo and GPU tensors through NCCL. This eliminates the need for manual device placement when performing allreduce operations on base types.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/thop/ncclWindowTensor.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/thop/ncclWindowTensor.htensorrt_llm/functional.pycpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-08-19T12:45:11.997Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.py
📚 Learning: 2025-08-26T06:07:02.166Z
Learnt from: shaharmor98
Repo: NVIDIA/TensorRT-LLM PR: 7231
File: tensorrt_llm/_torch/pyexecutor/_util.py:504-509
Timestamp: 2025-08-26T06:07:02.166Z
Learning: In tensorrt_llm/_torch/pyexecutor/_util.py, when calling model_engine.set_lora_model_config(), pass model_binding_config.mlp_hidden_size directly without multiplying by mapping.tp_size, as the mlp_hidden_size from get_bindings_model_config() is already the per-TP rank value needed for LoRA weight packaging.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.py
📚 Learning: 2025-09-22T19:25:45.607Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cpp:170-179
Timestamp: 2025-09-22T19:25:45.607Z
Learning: In NCCLUserBufferAllocator::getNCCLDevComm(), multimem support is hard-coded to true because multimem is required for this function. The caller is responsible for ensuring multimem is available before calling this function - it should not be called if multimem is not supported.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.pycpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cppcpp/tensorrt_llm/thop/allreduceOp.cppcpp/tensorrt_llm/common/ncclUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-09-02T13:42:44.885Z
Learnt from: pcastonguay
Repo: NVIDIA/TensorRT-LLM PR: 7455
File: tensorrt_llm/_torch/pyexecutor/py_executor.py:1852-1860
Timestamp: 2025-09-02T13:42:44.885Z
Learning: In MPI communication within TensorRT-LLM pipeline parallelism, different communication types (tokens, logits, termination sync) must use disjoint tag namespaces to avoid message routing collisions when using the same source/destination patterns.
Applied to files:
cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp
📚 Learning: 2025-08-15T06:46:54.897Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.
Applied to files:
cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/common/ncclUtils.h
📚 Learning: 2025-08-08T05:10:38.906Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:0-0
Timestamp: 2025-08-08T05:10:38.906Z
Learning: The ScaledAccPerRowBiasPerColScaleScatter fusion in CUTLASS extensions (cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp) is specifically designed for per-column scaling factors only, so it uses a fixed Stride<_0,_1,int64_t> rather than conditional stride logic.
Applied to files:
cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cppcpp/tensorrt_llm/common/customAllReduceUtils.hcpp/tensorrt_llm/thop/allreduceOp.cpp
📚 Learning: 2025-09-23T15:13:48.819Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/multimem.h:20-30
Timestamp: 2025-09-23T15:13:48.819Z
Learning: TRT-LLM targets modern CUDA toolkits that support FP8 datatypes, so cuda_fp8.h can be included unconditionally without version guards in TRT-LLM code.
Applied to files:
cpp/tensorrt_llm/common/opUtils.cppcpp/tensorrt_llm/common/customAllReduceUtils.hcpp/tensorrt_llm/thop/allreduceOp.cpp
📚 Learning: 2025-08-21T21:48:35.135Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:399-417
Timestamp: 2025-08-21T21:48:35.135Z
Learning: CUTLASS extensions in TensorRT-LLM (located under cpp/tensorrt_llm/cutlass_extensions/) are designed to integrate with and extend functionality in the external CUTLASS repository. When analyzing these extensions, their consumers and functionality wiring may exist in the CUTLASS codebase rather than within TensorRT-LLM itself.
Applied to files:
cpp/tensorrt_llm/common/opUtils.cpp
📚 Learning: 2025-08-21T09:41:49.347Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.
Applied to files:
cpp/tensorrt_llm/common/opUtils.cpp
📚 Learning: 2025-08-08T05:06:31.596Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:36-36
Timestamp: 2025-08-08T05:06:31.596Z
Learning: CUTLASS extension files (under cpp/tensorrt_llm/cutlass_extensions/) follow CUTLASS coding style conventions, including using #pragma once instead of TRTLLM_ prefixed header guards, even though they are .hpp files.
Applied to files:
cpp/tensorrt_llm/thop/ncclWindowTensor.h
📚 Learning: 2025-08-14T15:38:01.771Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Applied to files:
tensorrt_llm/functional.py
📚 Learning: 2025-07-28T17:06:08.621Z
Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
Applied to files:
tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py
📚 Learning: 2025-09-09T09:40:45.658Z
Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.
Applied to files:
tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py
📚 Learning: 2025-08-20T06:56:02.889Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Applied to files:
cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp
📚 Learning: 2025-08-21T02:39:12.009Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.
Applied to files:
cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp
🧬 Code graph analysis (11)
tests/unittest/_torch/multi_gpu/test_allreduce.py (1)
tensorrt_llm/functional.py (1)
AllReduceStrategy(3876-3885)
cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp (2)
cpp/tensorrt_llm/common/customAllReduceUtils.h (1)
getMaxRequiredWorkspaceSize(34-45)cpp/tensorrt_llm/thop/allreduceOp.cpp (1)
rank(873-984)
cpp/tensorrt_llm/thop/ncclWindowTensor.cpp (1)
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp (3)
comm(81-108)comm(208-241)comm(536-569)
cpp/tensorrt_llm/common/opUtils.cpp (1)
cpp/tensorrt_llm/common/ncclUtils.cpp (6)
getInstance(34-38)getInstance(34-34)getInstance(122-126)getInstance(122-122)getInstance(244-248)getInstance(244-244)
cpp/tensorrt_llm/thop/ncclWindowTensor.h (1)
cpp/tensorrt_llm/thop/ncclWindowTensor.cpp (2)
create_nccl_window_tensor(23-42)create_nccl_window_tensor(23-24)
tests/microbenchmarks/all_reduce.py (1)
tensorrt_llm/functional.py (1)
AllReduceStrategy(3876-3885)
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp (2)
cpp/tensorrt_llm/common/opUtils.cpp (2)
getComm(76-147)getComm(76-76)cpp/tensorrt_llm/common/ncclUtils.cpp (6)
getInstance(34-38)getInstance(34-34)getInstance(122-126)getInstance(122-122)getInstance(244-248)getInstance(244-244)
tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py (4)
tensorrt_llm/_torch/distributed/ops.py (1)
AllReduce(554-710)tensorrt_llm/functional.py (1)
AllReduceStrategy(3876-3885)tensorrt_llm/mapping.py (1)
Mapping(351-510)tests/unittest/conftest.py (1)
mpi_pool_executor(246-254)
cpp/tensorrt_llm/thop/allreduceOp.cpp (2)
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp (3)
comm(81-108)comm(208-241)comm(536-569)cpp/tensorrt_llm/common/ncclUtils.cpp (6)
getInstance(34-38)getInstance(34-34)getInstance(122-126)getInstance(122-122)getInstance(244-248)getInstance(244-244)
cpp/tensorrt_llm/common/ncclUtils.cpp (2)
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp (3)
comm(81-108)comm(208-241)comm(536-569)cpp/tensorrt_llm/common/ncclUtils.h (3)
NCCLHelper(56-167)isLoaded(72-143)NCCLWindowBuffer(198-300)
cpp/tensorrt_llm/common/ncclUtils.h (2)
cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.h (1)
tensorrt_llm(26-79)cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp (4)
comm(102-105)comm(102-102)ptr(22-25)ptr(22-22)
🪛 Clang (14.0.6)
cpp/tensorrt_llm/thop/ncclWindowTensor.h
[error] 18-18: 'tensorrt_llm/common/ncclUtils.h' file not found
(clang-diagnostic-error)
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp
[error] 17-17: 'tensorrt_llm/common/ncclUtils.h' file not found
(clang-diagnostic-error)
cpp/tensorrt_llm/common/ncclUtils.cpp
[error] 17-17: 'tensorrt_llm/common/ncclUtils.h' file not found
(clang-diagnostic-error)
cpp/tensorrt_llm/common/ncclUtils.h
[error] 18-18: 'tensorrt_llm/common/assert.h' file not found
(clang-diagnostic-error)
🪛 Ruff (0.14.5)
tests/unittest/_torch/multi_gpu/test_ncclwindowtensor.py
56-56: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
56-56: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
104-104: Unused function argument: tensor_parallel_rank
(ARG001)
145-145: Unused function argument: tensor_parallel_rank
(ARG001)
161-161: Loop control variable i not used within loop body
(B007)
190-190: Unused function argument: tensor_parallel_rank
(ARG001)
226-226: Unused function argument: tensor_parallel_rank
(ARG001)
375-387: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
406-419: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
437-449: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
467-479: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
497-509: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (14)
cpp/tensorrt_llm/common/customAllReduceUtils.h (1)
146-155: LUT bounds-check fallback toNCCL_SYMMETRIClooks consistentThe extended bounds check before indexing
AllReduceBestStrategyTableand falling back toAllReduceStrategyType::NCCL_SYMMETRICfor out-of-range entries is safe and aligns with the PR’s new default symmetric strategy. The short-circuiting||chain correctly avoids anyat()calls whensm_versionis missing.tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py (1)
29-34: Enum mapping and default LUT initialization align with NCCL_SYMMETRIC fallbackAdding
'NCCL_SYMMETRIC': 8tostrategy_name_to_enumand initializingstrategy_tablewithConstants.strategy_name_to_enum['NCCL_SYMMETRIC']keeps the Python generator consistent with theAllReduceStrategyIntEnum (whereNCCL_SYMMETRIC == 8) and with the C++ fallback you added. When regenerating LUTs, missing or filtered entries will now correctly default to NCCL_SYMMETRIC.Also applies to: 88-92
tests/unittest/_torch/multi_gpu/test_allreduce.py (1)
114-131: Test now targets NCCL_SYMMETRIC path as intendedSwitching
Linear(..., allreduce_strategy=AllReduceStrategy.NCCL_SYMMETRIC)ensures the fusion tests exercise the new default NCCL_SYMMETRIC strategy instead of the legacy NCCL path, matching the rest of the PR’s behavior change. No issues spotted.tests/microbenchmarks/all_reduce.py (1)
243-249: Expanded pandas display options are fine for this CLI benchmarkThe additional
pd.set_optioncalls to show all columns and avoid wrapping are appropriate for human inspection of the results on rank 0 and remain scoped to this script.cpp/tensorrt_llm/thop/CMakeLists.txt (1)
37-107:ncclWindowTensor.cppcorrectly wired intoth_commonAdding
ncclWindowTensor.cppto theth_commonsources is the right place for exposing the new Torch op alongside the other thop kernels.cpp/tests/unit_tests/multi_gpu/CMakeLists.txt (1)
17-24: NewncclUtilsTesttarget wiring looks correctThe new
ncclUtilsTestgtest target and its conditional linkage to Python and Torch libraries underBUILD_PYTare consistent with the existing multi-GPU test setup.cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cpp (1)
23-27: SimplifiedUserBufferAllocator::Instance()singleton is appropriateReturning a single function-local static
UserBufferAllocatorinstance matches the removal of the NCCL-specific allocator path and keeps the API straightforward and thread-safe.cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp (1)
215-240: Fallback behavior updated cleanly to NCCL_SYMMETRICThe selection heuristics and logging for non‑P2P, non‑NVLINK, and oversized workspace cases now consistently fall back to
AllReduceStrategyType::NCCL_SYMMETRIC, with appropriate deterministic‑mode warnings. This looks correct and matches the high‑level goal of changing the default fallback from NCCL to NCCL_SYMMETRIC.Also applies to: 272-316
cpp/tensorrt_llm/thop/allreduceOp.cpp (2)
18-24: Verify header wiring forncclUtils.hIncluding
"tensorrt_llm/common/ncclUtils.h"here is expected given the new NCCL window utilities, but clang static analysis is reportingfile not found. Please confirm that:
ncclUtils.his added to the relevant target’s include paths, and- The CMake target that builds this file depends on the library exposing that header.
This may just be a stale compile database, but it’s worth checking the build configuration.
270-305: Runtime strategy dispatch including NCCL_SYMMETRIC looks consistentThe
run()method’s switch now cleanly dispatchesAllReduceStrategyType::NCCL_SYMMETRICtorunNCCLAllReduceSymmetric, leaving other strategies unchanged. This aligns with the updated heuristic selection and keeps UB/LOWPRECISION/fusion paths intact.cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp (1)
41-72: Good coverage of communicator resource lifecycle and window allocator behaviorThe
createSplitCommhelper and associated tests do a nice job of:
- Mirroring
getComm’s deleter pattern (resource cleanup viaNcclCommResourceManagerbeforencclCommDestroy),- Verifying registration, ordering, and count semantics of
NcclCommResourceManager, and- Exercising
NCCLWindowAllocatoracross basic allocation, reuse, best‑fit selection, scoped buffers, multi‑comm isolation, and cleanup on comm destruction, plus PyTorchcreateNCCLWindowTensorintegration when enabled.This should give good confidence that the new NCCL utility layer behaves correctly across realistic communicator lifecycles.
Also applies to: 121-199, 205-527
cpp/tensorrt_llm/common/ncclUtils.cpp (1)
244-293: NCCL window allocator implementation and cleanup logic look sound post‑fixAside from the symbol‑loading issues above, the core of
NCCLWindowAllocatorlooks solid:
requestBuffer:
- Uses a per‑comm best‑fit search to reuse the smallest adequate buffer.
- Registers a per‑comm cleanup callback only once via
NcclCommResourceManager.- Pool bookkeeping:
getBufferCountandgetBufferInUseCountcorrectly reflect pool vs in‑use counts.releaseBufferandScopedNCCLWindowBuffercooperate to return buffers to the pool.- Cleanup:
cleanupBuffersForCommsynchronizes the device, deregisters all windows viancclCommWindowDeregister, and frees device memory viancclMemFree, with warnings logged on errors and then erases all state for that comm.Given the unit tests in
cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp, this design should behave correctly under comm teardown and buffer reuse once the symbol‑loading guardrails are in place.Also applies to: 447-565
cpp/tensorrt_llm/common/ncclUtils.h (2)
54-96: Dynamic NCCL helper interface looks good and aligns with runtime symbol loading designThe
NCCLHelperinterface (singleton, typed function pointers forncclCommWindowRegisterandncclMemAlloc,isLoaded()query) matches the project’s pattern of deferring NCCL feature detection to runtime, and it fits well with the existing “NCCL library already loaded by PyTorch/bindings” assumption. The header side looks clean and self‑contained; correctness will hinge on the.cppimplementation ofloadNCCLLibraryand symbol resolution.
232-300: Window allocator cleanup robustness confirmed—no action requiredThe implementation correctly handles communicator reuse. The
cleanupBuffersForCommfunction erases the communicator from bothmBufferPool(line 562) andmRegisteredComms(line 563), ensuring that when NCCL reuses ancclComm_taddress, the old cleanup records are properly removed. TheregisterBufferCleanupguard check at line 450 prevents duplicate registrations for new communicators at reused addresses.
| { | ||
| // If messageSize is less than maxWorkspaceSize, use NCCL, regardless of the fusion type. | ||
| // If messageSize is greater than maxWorkspaceSize or topology is unsuitable, use NCCL_SYMMETRIC fallback. | ||
| if (message_size_bytes > max_workspace_size || !mIsP2PSupported || !mIsNVLINKSupported) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does nccl symmetric support PCIe too? Do we any benchmarks on PCIe or mixed PCIe NVLink systems to make sure we don't have any regressions?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, NCCL_SYMMETRIC should also work with PCIe.
But I assume the benefit is smaller.
I am finding a system with PCI and test that there as well.
That might take a little bit though.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I updated the code to detect NVLink and MNNVL and apply this optimization only if one of them is detect to avoid regression on PCIe systems.
querying MNNVL to determine if it is worth it to copy data
15373a9 to
429dcaa
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #25764 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #25772 [ run ] triggered by Bot. Commit: |
|
PR_Github #25764 [ run ] completed with state |
|
PR_Github #25772 [ run ] completed with state |
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
81e4909 to
937a373
Compare
|
PR_Github #25786 [ run ] triggered by Bot. Commit: |
|
PR_Github #25786 [ run ] completed with state |
|
/bot run --disable-fail-fast --reuse-test |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand.
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #25871 [ run ] triggered by Bot. Commit: |
…f the nccl communicator Signed-off-by: Ludwig Schneider <[email protected]> small fix Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]> removing debug output Signed-off-by: Ludwig Schneider <[email protected]> removing debug output Signed-off-by: Ludwig Schneider <[email protected]> slowing moving to use NCCL_SYMMETRIC more Signed-off-by: Ludwig Schneider <[email protected]> fixing output of all_reduce.py benchmark Signed-off-by: Ludwig Schneider <[email protected]> eliminate cudamemcopy for very small buffers Signed-off-by: Ludwig Schneider <[email protected]> remove stupid check Signed-off-by: Ludwig Schneider <[email protected]> modelling perf Signed-off-by: Ludwig Schneider <[email protected]> adding files for python bindings Signed-off-by: Ludwig Schneider <[email protected]> fixing namespace issue Signed-off-by: Ludwig Schneider <[email protected]> make it work with pytorch Signed-off-by: Ludwig Schneider <[email protected]> finding out what is going wrong Signed-off-by: Ludwig Schneider <[email protected]> simplification on the locking Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
Signed-off-by: Ludwig Schneider <[email protected]>
581621c to
f4a0f84
Compare
|
PR_Github #25871 [ run ] completed with state |
Summary by CodeRabbit
New Features
Bug Fixes
Deprecations
Tests
Description
Background and Motivation
Currently, when AllReduce operations encounter conditions that prevent the use of optimized strategies (e.g., large message sizes, missing P2P support, or out-of-bounds lookup table entries), the system falls back to the
NCCLstrategy. This PR introduces a new NCCL window tensor infrastructure and changes all default fallback paths to useNCCL_SYMMETRICinstead, which provides better performance characteristics through window tensor registration and improved buffer management.The
NCCL_SYMMETRICstrategy leverages NCCL window tensors for efficient buffer reuse and registration, making it a more suitable default fallback than the basicNCCLstrategy. This PR includes the infrastructure necessary to support window tensor operations.Summary of Changes
This PR consists of two main components:
1. New NCCL Window Tensor Infrastructure
Introduces a new NCCL utilities system (
ncclUtils.cpp/h) that provides:NCCLWindowAllocator: Manages NCCL window-registered buffers with pooling and automatic cleanup. Buffers are tied to the lifetime of their associated NCCL communicator, enabling efficient buffer reuse across multiple AllReduce operations.NcclCommResourceManager: Thread-safe singleton that manages resources associated with NCCL communicators. Ensures proper cleanup of window buffers and other resources before communicator destruction.NCCLHelper: Dynamic library loading for NCCL symbols (ncclCommWindowRegister,ncclMemAlloc), allowing graceful handling of NCCL versions with or without window support.createNCCLWindowTensor: Helper function to create PyTorch tensors backed by NCCL window-registered buffers.This infrastructure decouples the NCCL Window allocation from the UB tensor allocation mechanism.
2. Default Fallback Strategy Changes
Systematically updates all fallback paths in the AllReduce strategy selection logic to use
NCCL_SYMMETRICinstead ofNCCL:allreduceOp.cpp,allreducePlugin.cpp):selectImplementation()methods to returnNCCL_SYMMETRICfor all fallback conditionsrunNCCLAllReduceSymmetric()now uses the newNCCLWindowAllocatorinstead of UB allocatorcustomAllReduceUtils.h): Updated lookup table fallback andSelectStrategyLP()to returnNCCL_SYMMETRICfunctional.py): Updated workspace fallback logic to useNCCL_SYMMETRICallreduce_heuristic_code_gen.py): Updated default lookup table initialization to useNCCL_SYMMETRICncclUtilsTest.cpp) for the new infrastructure:NCCLWindowAllocatorbuffer allocation, reuse, and cleanupNcclCommResourceManagerresource registration and cleanuptest_window_tensor.py) covering:Files Changed
New Infrastructure
cpp/tensorrt_llm/common/ncclUtils.h- NEW header for NCCL utilities (window allocator, resource manager, helper)cpp/tensorrt_llm/common/ncclUtils.cpp- NEW implementation of NCCL utilitiescpp/tensorrt_llm/thop/ncclWindowTensor.h- NEW header for PyTorch window tensor creationcpp/tensorrt_llm/thop/ncclWindowTensor.cpp- NEW implementation of window tensor creationCore Implementation Updates
cpp/tensorrt_llm/thop/allreduceOp.cpp- Updated fallback returns, integratedNCCLWindowAllocator, replaced UB allocatorcpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp- Updated plugin fallback logic and loggingcpp/tensorrt_llm/common/customAllReduceUtils.h- Updated lookup table fallbacktensorrt_llm/functional.py- Updated Python fallback logictests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py- Updated default initializationPerformance Impact
NCCL_SYMMETRICuses window tensor registration for better buffer reuse, potentially improving performance for repeated AllReduce operationsNCCLWindowAllocatorprovides efficient buffer pooling with best-fit allocation strategyNCCLNCCL_SYMMETRICis already a tested and supported strategy, and the new infrastructure includes comprehensive testsFunctional Impact
NCCLtoNCCL_SYMMETRICNCCL_SYMMETRICnow uses NCCL window tensors instead of UserBuffers allocatorTLLM_NCCL_MIN_REGISTRATIONenv var)NCCLstrategy are unaffected; only AUTO mode and fallback paths are changedTest Coverage
C++ Unit Tests (
ncclUtilsTest.cpp)NCCLWindowAllocatorTests:BasicAllocation- Verifies buffer allocation and registrationBufferReuse- Tests buffer pooling and reuseBestFitReuse- Verifies best-fit allocation strategyMultipleBuffers- Tests concurrent buffer managementBufferSearch- Verifies buffer lookup functionalityCleanupOnCommDestroy- Ensures proper cleanup when communicator is destroyedMultiCommIsolation- Verifies buffers are isolated per communicatorNcclCommResourceManagerTests:ResourceRegistration- Tests resource registration and countingResourceCleanup- Verifies cleanup order and exception handlingThreadSafety- Tests concurrent resource registrationMultiCommResources- Tests resource isolation per communicatorExisting Tests
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.