tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py (1)

1-1: Add required NVIDIA Apache-2.0 header.

Per repository guidelines, prepend the 2025 NVIDIA Apache-2.0 header to this file.

Apply this at file top:

+# Copyright (c) 2025, NVIDIA CORPORATION. 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.

tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py (7)

1-1: Add required NVIDIA Apache-2.0 header.

Please prepend the 2025 NVIDIA Apache-2.0 header.

+# Copyright (c) 2025, NVIDIA CORPORATION. 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.

---

19-21: De-duplicate the eps literal via a module constant.

0.1253 appears in multiple places; define once to keep pattern, replacement, and registration in sync.

-def _allreduce_residual_rmsnorm_pattern(
-    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 0.1253
-):
+RMSNORM_EPS_DEFAULT = 0.1253
+
+def _allreduce_residual_rmsnorm_pattern(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = RMSNORM_EPS_DEFAULT
+):

---

30-40: Use the op’s default overload explicitly for consistency.

Elsewhere in this PR we call .default; align here to avoid packet vs. overload mismatches in FX.

-    hidden_states = torch.ops.auto_deploy.torch_dist_all_reduce(x)
+    hidden_states = torch.ops.auto_deploy.torch_dist_all_reduce.default(x)

---

60-61: OK to instantiate a fresh matcher.

If this transform can run multiple times in a pipeline, consider guarding duplicate registrations (see next comment).

---

63-79: Harden registration: avoid dupe patterns and centralize eps.

• Pass skip_duplicates=True to avoid re-registering on repeated transform runs.
• Reuse RMSNORM_EPS_DEFAULT in dummy args and scalar_workaround.

-        dummy_args = [
+        dummy_args = [
             torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16),  # x
             torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16),  # residual
             torch.randn(hidden, device="meta", dtype=torch.bfloat16),  # weight
-            0.1253,  # eps
+            RMSNORM_EPS_DEFAULT,  # eps
         ]
@@
-        register_ad_pattern(
+        register_ad_pattern(
             search_fn=_allreduce_residual_rmsnorm_pattern,
             replace_fn=_allreduce_residual_rmsnorm_repl,
             patterns=patterns,
             dummy_args=dummy_args,
             op_ignore_types={torch.ops.aten.to.dtype: (torch.dtype,)},
-            scalar_workaround={"eps": 0.1253},
+            scalar_workaround={"eps": RMSNORM_EPS_DEFAULT},
+            skip_duplicates=True,
         )

---

80-81: Recompile after graph edits to materialize codegen.

Pattern pass mutates gm.graph; call dead-code elimination and recompile.

-        num_matches = patterns.apply(gm.graph)
+        num_matches = patterns.apply(gm.graph)
+        gm.graph.eliminate_dead_code()
+        gm.recompile()

---

43-47: Wrap torch.ops.dist.fused_allreduce_residual_rmsnorm in try/except to fall back to the Python implementation
This avoids hard failures on builds without the fused op:

 try:
     _FUSED_ARRN = torch.ops.dist.fused_allreduce_residual_rmsnorm
 except AttributeError:
     from ...distributed.trtllm import fused_allreduce_residual_rmsnorm as _FUSED_ARRN

 def _allreduce_residual_rmsnorm_repl(x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float):
     return _FUSED_ARRN(x, residual, weight, eps)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.