[None][feat] Optimize nemotron-h from python level

[Wanli-Jiang]

• Enable more c++ routing combinations.
• Update mamba tensor operations.

Summary by CodeRabbit

• Refactor
  • Simplified expert routing logic for improved performance in certain configurations.
  • Optimized Mamba2 decoder path with persistent tensor pre-computation for faster inference execution.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

Two changes optimize tensor-handling logic: one simplifies routing eligibility conditions in a fused MOE implementation by removing an edge-case constraint, and another optimizes the Mamba2 decode path by pre-computing and caching expanded tensors rather than recomputing them per call.

Changes

Cohort / File(s)	Summary
Routing Eligibility Simplification tensorrt_llm/_torch/modules/fused_moe/routing.py	Removed topk_group == 1 and self.top_k != 22 clause from Deepseekv3RoutingImpl.noaux_tc logic; fused execution now disallowed only on num_experts > 512 or top_k > 8 and top_k != 22.
Mamba2 Decode Path Optimization tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py	Added persistent pre-expanded tensors (_A_expanded, _dt_bias_expanded, _D_expanded) computed once in post_load_weights and reused across decode calls; removed .contiguous() calls on rearranged tensors to preserve view semantics.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	The PR description provides two key bullet points describing the changes, but lacks detail in the Description and Test Coverage sections of the template. The PR Checklist is marked complete without evidence.	Expand the Description section to explain why these changes optimize nemotron-h and detail the test coverage for routing logic and mamba tensor operations.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[None][feat] Optimize nemotron-h from python level' is directly related to the changes made in the PR, which optimize nemotron-h configurations by simplifying routing logic and updating mamba tensor operations.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`:
- Around line 243-251: post_load_weights currently assigns plain tensor
attributes (_A_expanded, _dt_bias_expanded, _D_expanded) which are not moved
with module.to()/cuda(); replace those plain attributes with registered buffers
using self.register_buffer("<name>", tensor, persistent=False) so they follow
module device/dtype moves—i.e., register _A_expanded (ensure
.to(dtype=torch.float32) as before), _dt_bias_expanded, and _D_expanded after
creating them with repeat, using the same names so existing code that references
self._A_expanded, self._dt_bias_expanded, and self._D_expanded continues to
work.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: f424cb9b-0111-4ba8-b308-63bc342fce7d

📥 Commits

Reviewing files that changed from the base of the PR and between f2f9051 and f9beed0.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/modules/fused_moe/routing.py
• tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

[Wanli-Jiang]

/bot run --disable-fail-fast

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43399 [ run ] triggered by Bot. Commit: 5f969ea Link to invocation

[tensorrt-cicd]

PR_Github #43399 [ run ] completed with state SUCCESS. Commit: 5f969ea
/LLM/main/L0_MergeRequest_PR pipeline #33934 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43610 [ run ] triggered by Bot. Commit: a2b418b Link to invocation

[tensorrt-cicd]

PR_Github #43610 [ run ] completed with state SUCCESS. Commit: a2b418b
/LLM/main/L0_MergeRequest_PR pipeline #34103 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[yechank-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #43730 [ run ] triggered by Bot. Commit: a2b418b Link to invocation

[tensorrt-cicd]

PR_Github #43730 [ run ] completed with state FAILURE. Commit: a2b418b
/LLM/main/L0_MergeRequest_PR pipeline #34212 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43957 [ run ] triggered by Bot. Commit: 42d00cb Link to invocation

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #44089 [ run ] triggered by Bot. Commit: b40951e Link to invocation

[tensorrt-cicd]

PR_Github #44089 [ run ] completed with state SUCCESS. Commit: b40951e
/LLM/main/L0_MergeRequest_PR pipeline #34518 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #44245 [ run ] triggered by Bot. Commit: b40951e Link to invocation

[tensorrt-cicd]

PR_Github #44245 [ run ] completed with state FAILURE. Commit: b40951e
/LLM/main/L0_MergeRequest_PR pipeline #34667 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Wanli-Jiang]

/bot skip --comment "Tests were passed within two CI tests"

[tensorrt-cicd]

PR_Github #44334 [ skip ] triggered by Bot. Commit: b40951e Link to invocation

[tensorrt-cicd]

PR_Github #44334 [ skip ] completed with state SUCCESS. Commit: b40951e
Skipping testing for commit b40951e

Link to invocation