tensorrt_llm/_torch/auto_deploy/custom_ops/mxfp4.py (4)

1-3: Add NVIDIA copyright header.

All source files must include the current-year NVIDIA header.

Apply:

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+
 import torch

---

45-69: Docstring return mismatch.

Docstring says the op returns routed_out and router_logits, but the function returns only routed_out. Fix the docstring to avoid confusion in downstream users and tooling.

Apply:

-    Return:
-      routed_out:   same leading shape as hidden_states, last dim = H
-      router_logits: [T, E] (T = number of tokens = prod(hidden_states.shape[:-1]))
+    Return:
+      routed_out: same leading shape as hidden_states, last dim = H

---

6-13: Harden kernel hub import with a clear error message.

If kernels hub is unavailable, the import will fail at runtime without context. Raise a descriptive error.

Apply:

 def _hub():
     global _triton_kernels_hub
     if _triton_kernels_hub is None:
-        from kernels import get_kernel
-
-        _triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
+        try:
+            from kernels import get_kernel
+            _triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
+        except Exception as e:
+            raise RuntimeError(
+                "Failed to load Triton kernels hub 'kernels-community/triton_kernels'. "
+                "Ensure kernels are installed and discoverable."
+            ) from e
     return _triton_kernels_hub

---

15-26: Follow codebase globals convention (G_ prefix, UPPER_SNAKE_CASE).

_global _triton_kernels_hub violates the repository’s Python globals naming rule.

Apply:

-_triton_kernels_hub = None
+G_TRITON_KERNELS_HUB = None
@@
-    global _triton_kernels_hub
-    if _triton_kernels_hub is None:
+    global G_TRITON_KERNELS_HUB
+    if G_TRITON_KERNELS_HUB is None:
@@
-        _triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
-    return _triton_kernels_hub
+        G_TRITON_KERNELS_HUB = get_kernel("kernels-community/triton_kernels")
+    return G_TRITON_KERNELS_HUB

Also update its single reference at Line 17 accordingly.

tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4.py (5)

1-3: Add NVIDIA copyright header.

New source files must include the current-year NVIDIA header.

Apply:

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+
 from typing import Tuple

---

124-139: Router pattern return annotation is incorrect.

Function returns a single Tensor, but type hints declare a Tuple. Align annotation to avoid confusion and tooling issues.

Apply:

-def _router_pattern(
+def _router_pattern(
     hidden_states: torch.Tensor,
     weight: torch.Tensor,
     bias: torch.Tensor,
     top_k: int = 2,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> torch.Tensor:
@@
-    return router_scores
+    return router_scores

(Only the return type changes.)

---

150-163: Remove stray debug print.

Printing the FX graph during transforms is noisy and pollutes logs. Use logging if needed.

Apply:

-        print(str(graph))

---

215-218: Register expert blocks/scales as buffers, not parameters.

These tensors are inference-only metadata. Registering as parameters can perturb optimizers and state dict semantics.

Apply:

-def _ensure_param(mod: nn.Module, name: str, tensor: torch.Tensor) -> None:
-    if not hasattr(mod, name):
-        mod.register_parameter(name, nn.Parameter(tensor, requires_grad=False))
+def _ensure_param(mod: nn.Module, name: str, tensor: torch.Tensor) -> None:
+    if not hasattr(mod, name):
+        mod.register_buffer(name, tensor, persistent=True)

No other call sites need changes since attributes are accessed via get_attr nodes.

---

356-359: Potential name collisions for newly registered attributes.

gate_up_proj_blocks/scales and down_proj_blocks/scales are generic. If users re-run transforms, these may exist already. Consider suffixing with a stable tag or checking existence before registering.

You already guard with hasattr in _ensure_param, but add a prefix (e.g., ad) or a model-unique tag to avoid clashes across different experts modules if multiple MoE blocks exist.

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.