tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/rms_norm.py (1)

17-19: Consider expanding the docstring to follow Google-style format.

While the current docstring documents the key behavior, consider adding parameter descriptions to align with coding guidelines for Google-style docstrings.

Example expansion:

-    """Rms norm kernel.
-    Forces weights to be in float32 for the kernel.
-    """
+    """RMS normalization kernel.
+    
+    Applies RMS normalization with weights forced to float32 precision during
+    computation for numerical accuracy.
+    
+    Args:
+        output: Output tensor pointer.
+        input: Input tensor pointer.
+        weights: Weight tensor pointer.
+        eps: Epsilon value for numerical stability.
+        M: Number of rows.
+        N_COLS: Number of columns.
+        BLOCK_N: Block size for Triton kernel.
+    """

As per coding guidelines

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_rmsnorm.py (2)

29-29: Consider parameterizing the device.

The device is hard-coded to "cuda". While this is acceptable in a test file under singlegpu, consider accepting a device parameter for consistency with the TestModel pattern where other modules use device="cuda" inline.

---

37-37: Consider adding a comment about float32 weight precision.

The weight is explicitly cast to float32 during computation even though the module may be stored in float16. Adding an inline comment would clarify this intentional design choice.

For example:

-        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+        # Use float32 for weights to match Nemotron-style RMSNorm computation
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Learnt from: nvchenghaoz
PR: NVIDIA/TensorRT-LLM#8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.

tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/triton_moe.py (1)

253-258: EM upper-bound is safe; clarify assumption and guard when not capturing

• The bound EM = min(max_padded, Mtop_kBLOCK_SIZE_M) is a valid upper bound (actual padded tokens ≤ T*BLOCK_SIZE_M), so it won’t skip real work. The current comment about “top_ids … unique” is misleading.
• Suggest clearer comment and (optionally) a debug assert when not under CUDA graph capture.

Apply this diff to clarify and add a guarded assert:

-    if A.size(0) < config["BLOCK_SIZE_M"]:
-        # optimize for small batch_size.
-        # We assume that top_ids of each token is unique,
-        # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
-        # and we can skip some invalid blocks.
-        EM = min(sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"])
+    if A.size(0) < config["BLOCK_SIZE_M"]:
+        # Optimize for small batches by shrinking the launch grid.
+        # Let T = M * top_k. After per‑expert padding, the total tokens processed
+        # EM_actual ≤ T * BLOCK_SIZE_M and ≤ sorted_token_ids.size(0).
+        # Using EM = min(sorted_token_ids.size(0), T * BLOCK_SIZE_M) preserves correctness
+        # while cutting empty blocks when E » M.
+        EM = min(sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"])
+        # Optional safety check (avoid host sync during CUDA graph capture):
+        # if not torch.cuda.is_current_stream_capturing():
+        #     EM_actual = int(num_tokens_post_padded.item())
+        #     assert EM >= EM_actual

Please run a quick A/B to confirm numerical parity with and without this gate for small M (e.g., M ∈ {1, 2, 8}, various top_k and E).

tensorrt_llm/_torch/auto_deploy/config/default.yaml (1)

127-127: Backend switched to Triton: ensure env/op availability

Switching fuse_rmsnorm to backend: triton is fine; please confirm torch.ops.auto_deploy.triton_rms_norm is registered across all build variants and CI images. Also update any user docs referencing FlashInfer as default.

tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/rms_norm.py (2)

17-19: Docstring good; minor clarity

Docstring now reflects f32 weights. Consider noting output is cast back to input dtype.

---

31-31: Cast weight once; ensure eps in f32

Tiny cleanup: cast w to f32 at load and multiply directly; also make eps explicitly f32 to avoid mixed-precision surprises.

-    out = (w.to(tl.float32) * out).to(x.dtype)
+    w = w.to(tl.float32)
+    out = (w * out).to(x.dtype)

Optionally:

-    out = xf / tl.sqrt(var + eps)
+    out = xf / tl.sqrt(var + tl.full((), eps, tl.float32))

tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py (2)

155-168: Avoid duplicate registrations and stale docstring

• Add skip_duplicates=True to prevent redundant pattern registrations across dtype variants.
• Top-of-file docstring still says “using FlashInfer”; update to mention Triton.

-        search_fns = [
+        search_fns = [
             _rms_norm_pattern,
             _rms_norm_pattern_float32_weights,
         ]
         for search_fn in search_fns:
             for input_dtype, weight_dtype in configs:
                 register_ad_pattern(
                     search_fn=search_fn,
                     replace_fn=partial(_rms_norm_replacement, backend=self.config.backend),
                     patterns=patterns,
                     dummy_args=dummy_args(input_dtype, weight_dtype),
-                    op_ignore_types={},
+                    op_ignore_types={},
                     scalar_workaround={"eps": 1e-6},
+                    skip_duplicates=True,
                 )

---

1-1: Docstring mentions only FlashInfer

Update to “FlashInfer or Triton” to reflect current behavior.

-"""Graph transform to optimize RMSNorm execution using FlashInfer."""
+"""Graph transform to optimize RMSNorm execution using FlashInfer or Triton."""

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_rmsnorm.py (1)

89-99: Also parametrize Triton on the baseline path

Add (“triton”, torch.ops.auto_deploy.triton_rms_norm) to ensure non-Nemotron path is covered with the new default backend.

 @pytest.mark.parametrize("eps", [1e-2, 1e-6])
 @pytest.mark.parametrize(
     "variant, op",
     [
         ("flashinfer", torch.ops.auto_deploy.flashinfer_rms_norm),
         ("torch", torch.ops.auto_deploy.torch_rmsnorm),
+        ("triton", torch.ops.auto_deploy.triton_rms_norm),
     ],
 )
 def test_rmsnorm_fusion(eps, variant, op):

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Learnt from: nvchenghaoz
PR: NVIDIA/TensorRT-LLM#8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.