[TRTLLM-10004][feat] Enable GEMM -> AR with GEMM output in registered buffers#11589
Conversation
📝 WalkthroughWalkthroughReplaces per-call boolean Changes
Sequence Diagram(s)sequenceDiagram
participant Py as Python caller
participant TorchOps as torch.ops.trtllm
participant Cpp as C++ allocate_output
participant NCCL as NCCL util / comm
Py->>TorchOps: request kernel with output_buffer_kind, group
TorchOps->>Cpp: call allocate_output(output_size, dtype, device, BufferKind, group)
Cpp->>NCCL: if BufferKind==NcclWindow -> getComm(group) and createNCCLWindowTensor
alt comm & alloc success
NCCL-->>Cpp: window tensor (defined)
else fallback or disabled
NCCL-->>Cpp: failure / none
Cpp-->>Cpp: fallback to userbuffers or default empty tensor
end
Cpp-->>TorchOps: return allocated tensor
TorchOps-->>Py: kernel returns tensor (and validity flag for window creation where applicable)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (25)
tensorrt_llm/_torch/modules/gated_mlp.py (1)
158-160:⚠️ Potential issue | 🟠 Major
allow_window_outputis silently dropped in the LoRA dispatch path.
forward()acceptsallow_window_output(line 154) but the call toforward_loraon lines 159–160 does not pass it. As a result, whenever LoRA is active the NCCL window output optimisation is always disabled regardless of what the caller requested, while the non-LoRA path (line 166) correctly forwards it.🐛 Proposed fix
if bool(lora_params): - return self.forward_lora(x, all_rank_num_tokens, - final_all_reduce_params, lora_params) + return self.forward_lora(x, all_rank_num_tokens, + final_all_reduce_params, + allow_window_output=allow_window_output, + lora_params=lora_params)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/gated_mlp.py` around lines 158 - 160, The LoRA dispatch in forward() drops the allow_window_output flag; update the call to forward_lora inside forward (the branch that checks lora_params) to pass the allow_window_output parameter through (and if needed, adjust forward_lora's signature to accept this flag) so LoRA path honors the caller's allow_window_output setting just like the non-LoRA branch that calls the normal forward implementation.tensorrt_llm/_torch/autotuner.py (1)
1-3:⚠️ Potential issue | 🟡 MinorAdd NVIDIA Apache 2.0 header for this modified Python source
This file is missing the required NVIDIA copyright header after modification.
As per coding guidelines, "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/autotuner.py` around lines 1 - 3, The file tensorrt_llm._torch.autotuner (the module starting with the import block: ast, contextlib, copy) is missing the required NVIDIA Apache-2.0 header; add the standard NVIDIA Apache License, Version 2.0 header comment at the very top of the file including the correct copyright year of the latest meaningful modification and the NVIDIA copyright holder line, before any imports or code, so the file complies with the project's licensing header requirements.tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (2)
74-95:⚠️ Potential issue | 🟡 MinorPreserve module namespace in the lazy routing import
The current import pulls symbols directly; guidelines require keeping the module namespace.
🔧 Suggested adjustment
- from tensorrt_llm._torch.modules.fused_moe.routing import ( - ROUTING_METHOD_TYPE_TO_CLASS, RoutingMethodType) + from tensorrt_llm._torch.modules.fused_moe import routing @@ - if routing_method_type == RoutingMethodType.DeepSeekV3: + if routing_method_type == routing.RoutingMethodType.DeepSeekV3: @@ - routing_method = ROUTING_METHOD_TYPE_TO_CLASS[routing_method_type]( + routing_method = routing.ROUTING_METHOD_TYPE_TO_CLASS[routing_method_type]( top_k=top_k, **routing_cls_kwargs)As per coding guidelines, "Always maintain the namespace when importing. Use
from package.subpackage import fooinstead offrom package.subpackage.foo import SomeClassorimport package."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py` around lines 74 - 95, The lazy import currently pulls symbols directly from tensorrt_llm._torch.modules.fused_moe.routing; replace that with a namespaced lazy import (e.g., import tensorrt_llm._torch.modules.fused_moe.routing as routing_module) and update usages: use routing_module.RoutingMethodType for the enum check and routing_module.ROUTING_METHOD_TYPE_TO_CLASS when constructing routing_method (and any other direct symbol references like ROUTING_METHOD_TYPE_TO_CLASS or RoutingMethodType) so the module namespace is preserved.
1-3:⚠️ Potential issue | 🟡 MinorAdd NVIDIA Apache 2.0 header for this modified Python source
This file is missing the required NVIDIA copyright header after modification.
As per coding guidelines, "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py` around lines 1 - 3, Add the required NVIDIA Apache-2.0 copyright header at the top of tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py: insert the standard multi-line NVIDIA Apache 2.0 license block (with the year of latest meaningful modification) immediately above the first import, so the header precedes the existing imports (dataclass, lru_cache, typing imports) and any module-level code; ensure the header matches the project's canonical Apache 2.0 header format and includes the correct year and copyright holder.cpp/tensorrt_llm/thop/allreduceOp.cpp (1)
1-4:⚠️ Potential issue | 🟡 MinorUpdate SPDX copyright year to 2026
This file was modified, but the SPDX header still ends at 2024.
🔧 Suggested update
- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &As per coding guidelines, "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/allreduceOp.cpp` around lines 1 - 4, The SPDX header at the top containing "SPDX-FileCopyrightText" currently lists the year range ending in 2024; update that year to 2026 (e.g., change "1993-2024" to "1993-2026") while preserving the rest of the comment and the "SPDX-License-Identifier: Apache-2.0" token so the file header reflects the latest modification year.tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py (2)
1-3:⚠️ Potential issue | 🟡 MinorAdd NVIDIA Apache 2.0 header for this modified Python source
This file is missing the required NVIDIA copyright header after modification.
As per coding guidelines, "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py` around lines 1 - 3, The file tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py is missing the required NVIDIA Apache-2.0 copyright header; add the standard NVIDIA Apache License 2.0 header block (including the year of latest meaningful modification) at the very top of the file above the existing imports (e.g., before "from operator import getitem") so the file contains the correct license and copyright notice.
528-575:⚠️ Potential issue | 🟠 MajorSearch pattern parameter mismatch prevents prologue optimization from triggering
The search pattern at line 535 binds
to_userbuffersas a keyword argument, but the actualnvfp4_gemmfunction signature usesoutput_buffer_kind: int. This mismatch means the pattern will never match any call sites, and the UB prologue replacement optimization will be silently skipped.Update the search pattern to bind
output_buffer_kindinstead:
- Change
to_userbuffers_key = KeywordArg('to_userbuffers')tooutput_buffer_kind_key = KeywordArg('output_buffer_kind')- Update the
CallFunctionto useoutput_buffer_kind=output_buffer_kind_keyinstead ofto_userbuffers=to_userbuffers_key- Update the
empty_nvfp4_gemm_prologue_patternsignature parameter fromto_userbuffers: booltooutput_buffer_kind: int🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py` around lines 528 - 575, The pattern binds the wrong keyword so the prologue never matches: in register_nvfp4_gemm_prologue replace the to_userbuffers binding with output_buffer_kind by renaming to_userbuffers_key = KeywordArg('to_userbuffers') to output_buffer_kind_key = KeywordArg('output_buffer_kind'), update the CallFunction trtllm_nvfp4_gemm_default to use output_buffer_kind=output_buffer_kind_key instead of to_userbuffers=..., and change the empty_nvfp4_gemm_prologue_pattern signature parameter from to_userbuffers: bool to output_buffer_kind: int so the search pattern matches the actual torch.ops.trtllm.nvfp4_gemm signature used in target_nvfp4_gemm_prologue_pattern.tensorrt_llm/_torch/distributed/ops.py (1)
1-5:⚠️ Potential issue | 🟠 MajorAdd the NVIDIA copyright header.
This Python file is modified but missing the required Apache 2.0 NVIDIA header.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header
+# +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/distributed/ops.py` around lines 1 - 5, This file (module tensorrt_llm._torch.distributed.ops) is missing the required NVIDIA Apache-2.0 copyright header; add the standard NVIDIA header block (including the Apache 2.0 license boilerplate and the year of latest meaningful modification) at the very top of the file before any imports (before the existing import math line), ensuring the header text matches project guidelines and includes the correct year.cpp/tensorrt_llm/thop/fp4Gemm.cpp (1)
1-15:⚠️ Potential issue | 🟠 MajorUpdate copyright year to include 2026.
The file is modified but still ends at 2023.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header update
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/fp4Gemm.cpp` around lines 1 - 15, Update the file header copyright block in fp4Gemm.cpp to reflect the latest modification year by changing the year range "2020-2023" to include 2026 (e.g., "2020-2026"); ensure the rest of the NVIDIA Apache License header text remains unchanged so the file complies with the project's copyright-year policy.cpp/tensorrt_llm/thop/cudaNvfp4MM.cpp (1)
1-16:⚠️ Potential issue | 🟠 MajorUpdate copyright year to 2026.
The file was modified but still shows 2025 in the SPDX header.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header update
- * SPDX-FileCopyrightText: Copyright (out) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (out) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/cudaNvfp4MM.cpp` around lines 1 - 16, Update the SPDX header year from 2025 to 2026 in the file by editing the SPDX-FileCopyrightText line (the top comment block) so it reads the updated year; ensure the SPDX-License-Identifier and surrounding license text remain unchanged and the header formatting is preserved.cpp/tensorrt_llm/thop/weightOnlyQuantGemm.cpp (1)
1-15:⚠️ Potential issue | 🟠 MajorUpdate the copyright year to reflect 2026 modifications.
The file was modified, but the header still shows 2025. Please bump the year to the latest meaningful modification year.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header update
- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/weightOnlyQuantGemm.cpp` around lines 1 - 15, Update the top-of-file NVIDIA copyright header in weightOnlyQuantGemm.cpp to reflect the latest modification year by changing "2025" to "2026" so the header shows the year of latest meaningful modification; locate the header comment block at the beginning of the file and replace the year token accordingly.tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (1)
1-2:⚠️ Potential issue | 🟠 MajorAdd the NVIDIA copyright header.
This file is modified but missing the required Apache 2.0 NVIDIA header.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header
+# +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py` around lines 1 - 2, Add the required NVIDIA Apache-2.0 copyright header to the top of tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (above the existing imports), using the year of the latest meaningful modification; ensure the header matches the project's standard NVIDIA header text and license notice so the file complies with the coding guidelines.cpp/tensorrt_llm/thop/cudaScaledMM.cpp (1)
1-16:⚠️ Potential issue | 🟠 MajorUpdate copyright year range to include 2026.
The file is modified but the header ends at 2024.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header update
- * SPDX-FileCopyrightText: Copyright (out) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (out) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/cudaScaledMM.cpp` around lines 1 - 16, Update the file header copyright year range to include the latest modification year (2026); specifically edit the top comment block where the range "1993-2024" appears and change it to "1993-2026" so the SPDX header and license block reflect the current year.cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp (1)
1-15:⚠️ Potential issue | 🟠 MajorUpdate copyright year to include 2026.
The file is modified but still ends at 2023.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification.”📝 Suggested header update
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp` around lines 1 - 15, Update the copyright header in cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp to reflect the latest modification year by changing the year range that currently reads "2020-2023" to include 2026 (e.g., "2020-2026"); ensure the rest of the NVIDIA Apache 2.0 header text remains unchanged.tests/unittest/_torch/thop/parallel/test_fp4_linear.py (1)
1-1:⚠️ Potential issue | 🟡 MinorMissing NVIDIA copyright header on modified
.pyfile.The file has no SPDX copyright header at all. Per coding guidelines, all modified
.pysource files must carry an NVIDIA Apache-2.0 header with the year of latest meaningful modification.✏️ Proposed fix
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sysAs per coding guidelines, "Include NVIDIA copyright header on ALL new files and update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/thop/parallel/test_fp4_linear.py` at line 1, The file tests/unittest/_torch/thop/parallel/test_fp4_linear.py is missing the required NVIDIA Apache-2.0 copyright/SPDX header; add the standard NVIDIA Apache-2.0 header (including SPDX-License-Identifier: Apache-2.0), the NVIDIA copyright statement and the latest year of meaningful modification at the very top of the file before any imports (e.g., before the existing import sys), matching the project's header template exactly so all .py files carry the same license block.tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
1-1:⚠️ Potential issue | 🟠 MajorAdd NVIDIA Apache-2.0 header with updated year.
This modified file is missing the required NVIDIA copyright header (Apache License 2.0 format) with the latest modification year. Please add it above the existing content.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format.”🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/fused_moe/quantization.py` at line 1, Add the required NVIDIA Apache-2.0 copyright header (with the year of latest meaningful modification) at the very top of this file, above the existing content (before the import inspect line); ensure the header uses the standard Apache License 2.0 template used across the repo (including the copyright notice, SPDX identifier, and brief license text/notice block) and matches other files in style and wording.tensorrt_llm/_torch/models/modeling_deepseekv3.py (1)
1-1:⚠️ Potential issue | 🟠 MajorAdd NVIDIA Apache-2.0 header with updated year.
This modified file is missing the required NVIDIA copyright header (Apache License 2.0 format) with the latest modification year. Please add it above the existing content.
As per coding guidelines: “All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format.”🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_deepseekv3.py` at line 1, The file tensorrt_llm/_torch/models/modeling_deepseekv3.py is missing the required NVIDIA Apache-2.0 copyright header; add the NVIDIA Apache License, Version 2.0 header with the latest meaningful modification year at the very top of the file (above the existing leading line like "# --------------------------------------------------"), using the standard Apache-2.0 header text and including the NVIDIA copyright line and year.tensorrt_llm/_torch/models/modeling_nemotron_h.py (1)
1-1:⚠️ Potential issue | 🟡 MinorCopyright year should be updated for this modified file.
The header shows
2022-2024but this file is modified in the PR. As per coding guidelines, "update year on modified files."📝 Proposed fix
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_nemotron_h.py` at line 1, The SPDX copyright header at the top of modeling_nemotron_h.py still reads "2022-2024" and needs to reflect the current modification year; update that header line (the SPDX/FileCopyrightText comment) to include the current year (e.g., "2022-2026" or the appropriate end year) so the file header is accurate.cpp/tensorrt_llm/thop/cublasScaledMM.h (1)
1-2:⚠️ Potential issue | 🟡 MinorCopyright year needs updating on this modified file.
The header shows
2022-2024; file is modified in this PR. Per coding guidelines, "update year on modified files."📝 Proposed fix
-/* Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. +/* Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/cublasScaledMM.h` around lines 1 - 2, Update the file header copyright range in cpp/tensorrt_llm/thop/cublasScaledMM.h to include the current year (change "2022-2024" to "2022-2026"); edit the top-of-file comment block to reflect the new end year so the modified file complies with the project copyright-year policy.tensorrt_llm/_torch/modules/attention.py (1)
563-592:⚠️ Potential issue | 🟡 MinorDocument
allow_window_outputin the Args section.The new parameter is missing from the Google-style docstring, which makes the public API docs inaccurate.
✍️ Suggested docstring update
@@ attention_mask (AttentionMask): The attention mask type. mrope_config (Optional[dict]): The MROPE configuration. all_reduce_params (Optional[AllReduceParams]): The all reduce parameters. + allow_window_output (Optional[bool]): Whether to allow writing outputs into a preallocated NCCL window buffer. lora_params (Optional[dict]): The LoRA parameters. attention_window_size (Optional[int]): The attention window size. attention_mask_data (Optional[torch.Tensor]): The attention mask data.As per coding guidelines: "Use Google style docstrings for classes and functions, which can be parsed by Sphinx."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/attention.py` around lines 563 - 592, The forward method's docstring is missing the new allow_window_output parameter; update the Google-style Args section of Attention.forward to include allow_window_output (Optional[bool]) with a one-line description of its purpose and default/behavior (e.g., whether it enables windowed attention outputs when True and None meaning default behavior), ensuring the parameter name and type match the function signature and keeping the same docstring formatting as the other arguments.tensorrt_llm/_torch/models/modeling_llama.py (1)
555-563:⚠️ Potential issue | 🟠 Major
allow_window_outputbreaks MoE layers (unexpected kwarg).
Llama4DecoderLayernow passesallow_window_outputintoself.feed_forward, butLlama4MoE.forwarddoesn’t accept this parameter. This will raise aTypeErrorfor MoE layers. Please add the parameter and (optionally) use it to gate NCCL window allocation.🛠️ Suggested fix (signature + gating)
class Llama4MoE(nn.Module): @@ def forward( self, hidden_states: torch.Tensor, all_rank_num_tokens=None, final_all_reduce_params: Optional[AllReduceParams] = None, - cutlass_min_latency_mode: Optional[bool] = False, + cutlass_min_latency_mode: Optional[bool] = False, + allow_window_output: Optional[bool] = None, ) -> torch.Tensor: @@ - if not self.enable_attention_dp and self.mapping.has_tp(): + if allow_window_output and not self.enable_attention_dp and self.mapping.has_tp(): window = self.all_reduce.get_nccl_window_for_shape( shared_output.shape, all_reduce_params=final_all_reduce_params, like_tensor=shared_output, ) final_hidden_states = torch.add(shared_output, routed_output, out=window) + elif not self.enable_attention_dp and self.mapping.has_tp(): + final_hidden_states = shared_output + routed_output final_hidden_states = self.all_reduce( final_hidden_states, all_reduce_params=final_all_reduce_params) else: final_hidden_states = shared_output + routed_output🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_llama.py` around lines 555 - 563, The call in Llama4DecoderLayer passes allow_window_output into self.feed_forward but Llama4MoE.forward does not accept that kwarg, causing a TypeError; update the Llama4MoE.forward signature to accept allow_window_output: bool = False (or **kwargs) and use it to gate NCCL window allocation logic inside Llama4MoE (e.g., only allocate/use NCCL windows when allow_window_output is True), and ensure any call sites to feed_forward/MoE forwarding remain compatible by forwarding the flag through Llama4FeedForward/Llama4MoE entry points.tests/microbenchmarks/all_reduce.py (1)
33-93:⚠️ Potential issue | 🟡 MinorImport Linear via module namespace to match repo import rules.
Direct class imports violate the project’s namespace-preserving import style. Consider importing the module and referencing symbols from it (also avoids shadowing the local
linearvariable).✅ Suggested refactor
-from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode +from tensorrt_llm._torch.modules import linear as linear_module @@ - linear = None + linear_layer = None if profile_gemm_allreduce: @@ - linear = Linear( + linear_layer = linear_module.Linear( in_features=gemm_in_features, out_features=gemm_in_features, bias=False, dtype=input.dtype, mapping=mapping, - tensor_parallel_mode=TensorParallelMode.ROW, + tensor_parallel_mode=linear_module.TensorParallelMode.ROW, reduce_output=True, allreduce_strategy=strategy, ).to(input.device) - torch.nn.init.normal_(linear.weight, mean=0.0, std=0.02) + torch.nn.init.normal_(linear_layer.weight, mean=0.0, std=0.02) @@ - if profile_gemm_allreduce: - output = linear(x, all_reduce_params=allreduce_params) + if profile_gemm_allreduce: + output = linear_layer(x, all_reduce_params=allreduce_params)As per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/microbenchmarks/all_reduce.py` around lines 33 - 93, The file currently does "from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode" which violates the project's namespace-import rule and can shadow the local variable named "linear"; change to import the module (e.g., "import tensorrt_llm._torch.modules.linear as linear_mod") and update all references: replace "Linear" with "linear_mod.Linear" and "TensorParallelMode" with "linear_mod.TensorParallelMode" inside profile_allreduce (including the Linear(...) construction, the .to(...), and torch.nn.init.normal_(linear.weight) usage). Also consider renaming the local variable "linear" to "linear_layer" (and update its uses in func and the profile_gemm_allreduce branch) to avoid name shadowing with the module reference.cpp/tensorrt_llm/thop/cublasScaledMM.cpp (1)
2-2:⚠️ Potential issue | 🟡 MinorUpdate the header year to reflect the 2026 modification.
The file was modified, so the SPDX year should be current.
📄 Suggested header fix
- * SPDX-FileCopyrightText: Copyright (out) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (out) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines: "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/cublasScaledMM.cpp` at line 2, Update the SPDX header year to the year of latest modification (2026): locate the SPDX copyright/header block at the top of the file where the SPDX-FileCopyrightText and SPDX-License-Identifier are declared and change the year range "1993-2024" to "1993-2026" so the header reflects the current modification year.cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp (1)
2-2:⚠️ Potential issue | 🟡 MinorUpdate the NVIDIA copyright year to reflect the 2026 modification.
This file was modified, so the header year should be current.
📄 Suggested header fix
- * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.As per coding guidelines: "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp` at line 2, Update the file header copyright line to reflect the 2026 modification: replace the existing "Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved." line in the file header with the Apache License 2.0 style header showing the latest year "2026" and ensuring the header follows the same format used across the codebase (i.e., "Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved." and the standard Apache License 2.0 notice).tensorrt_llm/_torch/modules/linear.py (1)
389-402:⚠️ Potential issue | 🟡 Minor
output_buffer_kindandgroupare declared but unusedBoth parameters are never referenced in the method body. For
UnquantizedLinearMethod, the NCCL window path is handled viaapply_out; theapplysignature just adds noise and misleads callers/subclass authors.🔧 Proposed fix
- def apply(self, - module: Linear, - input: torch.Tensor, - bias: Optional[torch.Tensor], - output_buffer_kind: int = int(OutputBufferKind.DEFAULT), - group: Optional[List[int]] = None): + def apply(self, + module: Linear, + input: torch.Tensor, + bias: Optional[torch.Tensor]):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/linear.py` around lines 389 - 402, The apply method declares unused parameters output_buffer_kind and group which should be removed (or, if signature compatibility prevents removal, renamed to _output_buffer_kind and _group) to avoid misleading callers; update the UnquantizedLinearMethod.apply signature (the def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor], ...)) to drop those two params (or prefix them with underscores) and adjust any callers/overrides to match, noting that the NCCL/window handling is performed via apply_out so nothing needs to be done inside apply itself.
4e2a705 to
09e3f7f
Compare
|
/bot run --add-multi-gpu --disable-fail-fast |
|
PR_Github #36338 [ run ] triggered by Bot. Commit: |
|
PR_Github #36338 [ run ] completed with state
|
|
/bot run --add-multi-gpu --disable-fail-fast |
50ceac1 to
3886d90
Compare
|
PR_Github #36353 [ run ] triggered by Bot. Commit: |
|
PR_Github #36353 [ run ] completed with state
|
|
/bot run --add-multi-gpu --disable-fail-fast |
|
PR_Github #36519 [ run ] triggered by Bot. Commit: |
3886d90 to
8d9ac04
Compare
…indow_for_shape and prefer_window_output Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…nd from allocate_output Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…dow_output=True Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…+NVFP4 Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…lasses Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…mm limitation Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…alues Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…-GPU tests Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…cloudpickle traceback serialization Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…manager to fix cloudpickle serialization Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
…d e2e tests" This reverts commit 58b3293. Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
cloudpickle 3.1.2 with register_pickle_by_value serializes functions by value (embedding bytecode and globals). When the globals dict is serialized it re-encounters the function itself, creating a cycle that causes cloudpickle to attempt serializing torch.ops C++ objects (which are not picklable) and raises "TypeError: cannot pickle '_Ops' object". Fix: wrap direct torch.ops.trtllm.allocate_output() calls inside nested _allocate() helper functions in each _run_* helper. cloudpickle does not attempt to serialize nested function code objects as top-level by-value globals, so the _Ops reference is never encountered during pickling. This matches the pattern used in test_allreduce.py for the same reason. Also: move rank initialization before the try block (so rank=-1 sentinel is set on any failure path), tighten exception handling to BaseException (catches KeyboardInterrupt and SystemExit from MPI shutdown), and apply torch.inference_mode() as a context manager in run_single_rank rather than decorating the helpers (decorators wrap functions in callables that also fail cloudpickle by-value serialization). Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
274abe5 to
08422a7
Compare
|
/bot run --add-multi-gpu --disable-fail-fast |
|
PR_Github #45227 [ run ] triggered by Bot. Commit: |
|
PR_Github #45227 [ run ] completed with state
|
|
/bot run --add-multi-gpu --disable-fail-fast |
|
PR_Github #45409 [ run ] triggered by Bot. Commit: |
|
PR_Github #45409 [ run ] completed with state |
hyukn
left a comment
hyukn
left a comment
There was a problem hiding this comment.
Thanks for the great effort. LGTM.
… buffers (NVIDIA#11589) Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
10 participants
Summary by CodeRabbit
New Features
Refactor
Description
The goal is to use registered buffer for the NCCL_SYMMETRIC AllReduce strategy.
Output buffers are generated and not a problem, but input buffers must come from the previous operation directly into the registered buffer.
This PR implements this in fused python operations (GEMM and MOE) to enable output directly in a custom buffer if NCCL_SYMMETRIC is chosen.
This approach of modifying the python ops enables this to work independent of execution mode.
I.e. it works directly with default settings (torch compile is not required).
But as a downside, it manually needs to modify each operation individually that wants to make use of this speed up.
The FP4 GEMM can currently not be modified like this, since it already uses a differently defined custom output buffer.
(There might a future upside to unify this on a ncclWindowTensor as well.)
A little bit of plumbing is added to conservatively obtain the output tensors from the existing machinery.
I also modified the auto-tuner to use registered buffers for NCCL_SYMMETRIC assuming we can get those.
And the AR benchmark is modified to enable registered buffers as well.
Test Coverage
No explicit test for this is checked in, but literally any multi-GPU run is using this because of the auto-tuner.
PR Checklist
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse 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.