[None][feat] AutoDeploy: Use the router gemm op for nemotron MOE

[nvchenghaoz]

Summary by CodeRabbit

• Refactor
  • Improved routing computation efficiency through optimized data flow that conditionally selects between standard and specialized tensor operations, reducing overhead and enhancing model inference performance.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[nvchenghaoz]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

Router logits computation in NemotronHTopkRouter now conditionally branches on weight dtype: using standard linear operation for float32 weights, or a specialized Triton GEMM operator for other dtypes with transposed weights and float32 output forcing.

Changes

Cohort / File(s)	Change Summary
NemotronHTopkRouter Dtype-Conditional Routing tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py	Added dtype-based branching in forward method: float32 weights use standard F.linear, other dtypes invoke dsv3_router_gemm_op with transposed weight and forced float32 output dtype.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10–15 minutes

• Verify correctness of the Triton GEMM operator invocation and its compatibility with transposed weights
• Confirm dtype forcing to float32 doesn't cause unexpected precision loss or downstream issues
• Ensure both code paths produce equivalent or correctly-differentiated router logits

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is incomplete, containing only '@coderabbitai summary' without any explanation of the issue, solution, test coverage, or PR checklist completion.	Add a detailed description explaining the problem being solved, the solution implemented, relevant test coverage, and complete the PR checklist as specified in the template.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: using router gemm op for nemotron MOE in AutoDeploy, which matches the code changes that branched router logits computation based on weight dtype.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 356f67c and b687ad6.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., use from package.subpackage import foo and then foo.SomeClass() instead of from package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g., some_file.py)
Python class names should use PascalCase (e.g., class SomeClass)
Python function and method names should use snake_case (e.g., def my_awesome_function():)
Python local variable names should use snake_case, with prefix k for variable names that start with a number (e.g., k_99th_percentile = ...)
Python global variables should use upper snake_case with prefix G (e.g., G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g., self.x = 5 followed by """<type>: Description of 'x'""" )
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic

Files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

**/*.{cpp,h,cu,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top

Files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

🧠 Learnings (6)

📓 Common learnings

Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py:98-116
Timestamp: 2025-10-20T17:07:18.745Z
Learning: In NemotronH models (tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py), the gate (self.gate) returns topk_indices and topk_weights that are already in the correct shape to be passed directly to torch_ops.auto_deploy.torch_moe without needing to reshape them when hidden_states is flattened.

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 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.

📚 Learning: 2025-09-19T21:28:13.751Z

Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

📚 Learning: 2025-10-20T17:07:18.745Z

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py:98-116
Timestamp: 2025-10-20T17:07:18.745Z
Learning: In NemotronH models (tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py), the gate (self.gate) returns topk_indices and topk_weights that are already in the correct shape to be passed directly to torch_ops.auto_deploy.torch_moe without needing to reshape them when hidden_states is flattened.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

📚 Learning: 2025-10-20T17:09:21.560Z

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 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.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

📚 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/auto_deploy/models/patches/nemotron_h.py

📚 Learning: 2025-08-08T22:03:40.707Z

Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

⏰ 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 (1)

tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)

91-117: Router GEMM path looks consistent; just confirm supported dtypes/device for the custom op.

The new branching in _nemotron_h_topk_router_forward keeps shapes and final fp32 logits consistent across both paths:

• hidden_states.view(-1, self.config.hidden_size) feeds either:
  • F.linear(hidden_states.float(), self.weight) when self.weight.dtype is torch.float32, or
  • torch.ops.trtllm.dsv3_router_gemm_op(hidden_states, self.weight.t(), bias=None, out_dtype=torch.float32) otherwise.

This gives [N, num_experts] fp32 router_logits in both cases, which is exactly what noaux_tc_op expects and aligns with the pattern that downstream custom ops are fine consuming fp32 outputs.

The only thing worth double‑checking is that dsv3_router_gemm_op is registered and tested for all (hidden_states.dtype, self.weight.dtype) pairs you plan to use here (e.g., fp16/bf16 activations with fp8/bf16 weights) and that it is never hit on unsupported devices (like CPU) where only the F.linear path would have worked before. If there are any unsupported combos, you may want to tighten the condition or add a small assertion.

Please verify the supported dtype/device combinations for torch.ops.trtllm.dsv3_router_gemm_op in your environment and ensure that:

• All expected NemotronH router weights are either pure fp32 (taking the F.linear path) or use dtypes that the GEMM op supports.
• This code path is only exercised on devices where the trtllm router GEMM op is available (likely CUDA).

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions:
   | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context.
   Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[tensorrt-cicd]

PR_Github #25888 [ run ] triggered by Bot. Commit: b687ad6

[tensorrt-cicd]

PR_Github #25888 [ run ] completed with state SUCCESS. Commit: b687ad6
/LLM/main/L0_MergeRequest_PR pipeline #19632 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.