[FMDL-1328][feat] Add support for nano-v3 and super-v3 with pytorch backend

[Wanli-Jiang]

Features

• Verified with VANILLA and CUTLASS MoE backend.

• Support BF16 / FP8 / NVFP4 models.

• Support multi-stream for MoE shared and MoE chunking.

Summary by CodeRabbit

• New Features

  • Added Mixture-of-Experts support with flexible activation type configuration
  • Introduced support for Nemotron-Nano model variant

• Improvements

  • Enhanced weight quantization for MoE operations
  • Optimized parallel MoE execution with improved stream management

• Tests

  • Expanded test suite to cover additional Nemotron model variants

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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes add activation-type parameterization throughout the MoE quantization and weight-handling pipeline, introduce a new NemotronHMOE module with auxiliary CUDA stream support, extend weight mappers for MoE expert handling, and add utility functions for gated activation detection, weight splitting, and relu-squared computation. Changes span C++ kernels, Python model definitions, quantization logic, and test parameterization.

Changes

Cohort / File(s)	Summary
C++ MoE Quantization cpp/tensorrt_llm/thop/moeOp.cpp	Added base_activation_type parameter to FusedMoeRunner::getQuantParams(). Introduces expand_ratio derived from activation type to adjust weight validation sizes from fixed factor 2 to dynamic factors in MXFP4/MXF8 and NVFP4 branches.
HF Checkpoint Weight Mappers tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py, tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py	Updated import paths for split from local to centralized utils. Added MoE expert weight remapping logic in nemotron_h_weight_mapper to handle VANILLA backend (direct copy) and non-VANILLA backends (up_proj → w1/w3, down_proj → w2 with scale handling for FP8/NVFP4).
Nemotron-H Model Definition tensorrt_llm/_torch/models/modeling_nemotron_h.py	Introduced NemotronHMOE class implementing gated MoE with latent projection layers and auxiliary stream-based parallel execution. Extended NemotronHLayer to route layer type "E" to MoE and accept aux_stream_dict. Updated NemotronHModel to initialize auxiliary CUDA streams (MoeShared, MoeChunkingOverlap, MoeBalancer). Normalized rms_norm_eps in NemotronHForCausalLM from config.
MoE Module Factory & Interfaces tensorrt_llm/_torch/modules/fused_moe/create_moe.py, tensorrt_llm/_torch/modules/fused_moe/interface.py	Added activation_type parameter to create_moe() and propagated to backend constructors. Introduced internal is_gated_activation flag and intermediate_size_expand_ratio (2 for gated, 1 otherwise) in MoE base class for use in weight shape calculations.
MoE Backend Implementations tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py, tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py	Added activation_type and layer_idx parameters to both implementations. VanillaMoE includes gating-aware expert creation (MLP with relu2 for Relu2 activation, GatedMLP otherwise) and validation errors for unsupported non-gated configs. CutlassFusedMoE forwards activation type to base class and kernel.
MoE Quantization & Weight Handling tensorrt_llm/_torch/modules/fused_moe/quantization.py	Replaced hardcoded factor-2 multipliers with intermediate_size_expand_ratio in weight shape calculations (w3_w1_weight dimensions, w3_w1_weight_shape, scales). Updated split logic to use split_length = intermediate_size_per_partition * expand_ratio // 2 for w3/w1 slicing across FP8, NVFP4, and TRT variants.
Utility Functions tensorrt_llm/_torch/utils.py	Added is_gated_activation(ActivationType) → bool to identify Swiglu/SwigluBias/Geglu activations. Added split(x, tp_size, idx, dim=0) → torch.Tensor for tensor partitioning with divisibility validation. Added relu2(x) → torch.Tensor computing relu-squared via F.relu. Added import of torch.nn.functional as F.
Test Infrastructure tests/unittest/_torch/modeling/test_modeling_nemotron_h.py	Parameterized tests with model_folder to support Nemotron-H-8B-Base-8K and Nemotron-Nano-3-30B-A3.5B-dev-1024. Updated create_nemotron_h_llm() signature to accept and route model_folder for model path construction. Replaced static GPU memory skips with per-model conditional skips. Added model-specific reference logprobs, tolerances, and expectations (exact checks for smaller model, fuzzy comparison via similar() for larger model).

Sequence Diagram(s)

sequenceDiagram
    participant App as Application
    participant CreateMoE as create_moe()
    participant Backend as MoE Backend<br/>(Cutlass/Vanilla)
    participant Interface as MoE Interface
    participant Quantization as Quantization
    participant Kernel as Kernel/C++

    App->>CreateMoE: create_moe(..., activation_type)
    CreateMoE->>Backend: new Backend(..., activation_type)
    Backend->>Interface: super().__init__(..., activation_type)
    Interface->>Interface: is_gated = is_gated_activation(activation_type)
    Interface->>Interface: expand_ratio = 2 if is_gated else 1
    
    alt Vanilla MoE Path
        Backend->>Backend: if activation_type == Relu2<br/>create MLP experts
        Backend->>Backend: else create GatedMLP experts
    end
    
    alt Weight Loading Path
        Quantization->>Quantization: split_length = inter_size * expand_ratio // 2
        Quantization->>Quantization: allocate w3_w1 with expand_ratio scaling
        Quantization->>Kernel: pass expand_ratio to C++ quantParams
    end
    
    Backend->>Kernel: forward(..., activation_type)
    Kernel->>Kernel: getQuantParams(..., base_activation_type)<br/>adjust validation per activation type
    Kernel-->>Backend: result
    Backend-->>App: output

Loading

sequenceDiagram
    participant Model as NemotronHModel
    participant Layer as NemotronHLayer
    participant MoE as NemotronHMOE
    participant Router as Gate Router
    participant AuxStream as Aux CUDA Stream

    Model->>Model: __init__: create aux_stream_dict<br/>(MoeShared, Overlap, Balancer)
    Model->>Layer: pass aux_stream_dict
    
    Layer->>Layer: route layer_type=="E" to MoE
    Layer->>MoE: new NemotronHMOE(..., aux_stream_dict)
    MoE->>MoE: init latent projections (if enabled)
    MoE->>MoE: init gate and experts
    
    Layer->>MoE: forward(hidden_states)
    MoE->>Router: compute routing weights
    
    par Parallel Execution
        MoE->>MoE: shared path through gate
        MoE->>AuxStream: route to MoeShared stream
    and
        MoE->>MoE: expert path computation
        MoE->>AuxStream: route to MoeChunkingOverlap stream
    end
    
    AuxStream->>MoE: synchronize outputs
    MoE-->>Layer: combined result
    Layer-->>Model: propagate output

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

• C++ quantization logic (moeOp.cpp): New branching on expand_ratio requires verification against all quantization paths (MXFP4, NVFP8, NVFP4) to ensure size calculations remain correct and error messages align.
• New NemotronHMOE module (modeling_nemotron_h.py): Introduces untested parallel execution with auxiliary streams, latent projection logic, and new layer routing; requires careful verification of stream management and synchronization correctness.
• Activation-type propagation across multiple MoE backends: Dense, interconnected parameter threading through factory, interface, vanilla, and Cutlass implementations; each backend's gating-aware expert initialization requires independent reasoning.
• Weight remapping complexity (nemotron_h_weight_mapper.py): Non-trivial MoE expert weight transformation logic (up_proj → w1/w3 splitting, scale handling per backend) with multiple error paths that need coverage testing.
• Quantization weight shape updates (quantization.py): Widespread replacement of factor-2 with expand_ratio across multiple quantization variants (FP8, NVFP4, TRT) needs verification that slicing logic produces correct tensor dimensions for both gated and non-gated activations.
• Test parameterization (test_modeling_nemotron_h.py): Model-specific reference values and conditional tolerance logic; verify that each model's expected outputs and skip conditions are correctly mapped.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description includes a Features section listing key capabilities, but the required Description and Test Coverage sections are missing (only contain placeholder comments), and the PR checklist is incomplete.	Complete the Description section explaining what was changed and why, and the Test Coverage section listing relevant tests. Ensure all PR checklist items are properly addressed.
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.71% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main feature: adding support for nano-v3 and super-v3 models with the PyTorch backend, matching the changes throughout the pull request.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

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

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.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/tensorrt_llm/thop/moeOp.cpp (1)

538-542: runMoeMinLantency still assumes gated (×2) inter size, ignoring activation_type

runMoeMinLantency keeps this check:

TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0],
    "fc1_expert_weights and fc2_expert_weights must have the same number of experts.");
TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2,
    "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size.");

but you now support non-gated activations (e.g., ActivationType::Relu2) elsewhere by making the factor depend on isGatedActivation(base_activation_type).

For non-gated activations using the min-latency path, this will incorrectly require a 2× inter size and throw, even though the rest of the stack (workspace sizing, quant params) is activation-aware.

You can mirror runMoe’s logic here. One minimal change is:

-        TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0],
-            "fc1_expert_weights and fc2_expert_weights must have the same number of experts.");
-        TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2,
-            "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size.");
+        TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0],
+            "fc1_expert_weights and fc2_expert_weights must have the same number of experts.");
+
+        ActivationType base_activation_type = activation_type.has_value()
+            ? static_cast<ActivationType>(activation_type.value())
+            : ActivationType::Swiglu;
+        int expand_ratio = isGatedActivation(base_activation_type) ? 2 : 1;
+        TORCH_CHECK(
+            fc1_expert_weights.sizes()[1]
+                == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * expand_ratio,
+            expand_ratio == 2
+                ? "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."
+                : "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.");

and then drop the later re-declaration of base_activation_type near lines 556–558 (or reuse the same variable for activation params). That keeps min-latency MoE consistent with the main path for both gated and non-gated activations.

Also applies to: 556-559, 614-619

🧹 Nitpick comments (3)

tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py (1)

37-45: MoE expert remap logic looks correct; consider minor robustness tweaks

The new handling of mixer.{in,out}_proj *_scale and the MoE expert remap (mixer.experts.* -> w1/w2/w3) is consistent with standard gated-MoE layouts and NVFP4/FP8 scale formats. A couple of nits you may want to consider (not blockers):

• In the weight_scale and main-weight branches you rely on if weights[name].shape: to distinguish NVFP4 (tensor) vs FP8 (scalar). Making this explicit (e.g., checking weights[name].ndim == 0) would be clearer and less brittle.
• For the branches that slice weights[name].shape[0] // 2, an assert weights[name].shape[0] % 2 == 0 would defensively document the expectation that the combined dimension is 2×intermediate size.

Functionally this LGTM; the above are just clarity/defensiveness suggestions.

Also applies to: 98-130

tensorrt_llm/_torch/models/modeling_nemotron_h.py (2)

128-146: Consider per-layer handling for moe_intermediate_size lists

NemotronHMOE sets:

self.moe_intermediate_size = config.moe_intermediate_size[0] \
    if isinstance(config.moe_intermediate_size, list) else config.moe_intermediate_size

While MLPLayer treats list-valued intermediate_size as:

• Use intermediate_size[0] when len == 1 (shared across layers).
• Otherwise index by layer_idx.

If config.moe_intermediate_size can be a per-layer list (len > 1), always taking index 0 will give the wrong width for later MoE layers. You may want to mirror the MLPLayer pattern here, e.g.:

-        self.moe_intermediate_size = config.moe_intermediate_size[0] \
-            if isinstance(config.moe_intermediate_size, list) else config.moe_intermediate_size
+        if isinstance(config.moe_intermediate_size, list):
+            if len(config.moe_intermediate_size) == 1:
+                self.moe_intermediate_size = config.moe_intermediate_size[0]
+            else:
+                self.moe_intermediate_size = config.moe_intermediate_size[self.layer_idx]
+        else:
+            self.moe_intermediate_size = config.moe_intermediate_size

If current configs only ever use a scalar or a single-element list, behavior stays unchanged; this just makes the implementation future-proof for per-layer MoE widths.

---

147-221: NemotronHMOE MoE wiring and multi-stream usage look consistent

A few points on the new NemotronHMOE:

• Using ActivationType.Relu2 and passing it into create_moe(... activation_type=...) aligns with the new activation-type-aware MoE kernels.

• Gate (DeepseekV3Gate) and experts (create_moe with MoEWeightLoadingMode.VANILLA) are wired the same way as other DeepSeekV3-style MoE modules, and the latent projection path is correctly guarded by use_latent_moe.

• Pulling all_rank_num_tokens from attn_metadata and forwarding it into self.experts(..., all_rank_num_tokens=..., use_dp_padding=False) matches how MoE uses DP metadata elsewhere.

• The multi-stream call

  routed_output, shared_output = maybe_execute_in_parallel(
      _compute_routed_output, _compute_shared_output,
      self.event_dict[EventType.Main],
      self.event_dict[EventType.MoeShared], self.aux_stream_shared)

  is consistent with the helper’s contract and ensures both paths complete before you sum and reshape.

Only minor nit: forward(..., **kwargs) intentionally ignores extra kwargs (e.g., mamba_metadata) to stay signature-compatible with other mixers; if Ruff’s ARG002 is noisy you could rename to _kwargs, but it’s not functionally important.

Also applies to: 222-267

📜 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 ee941ac and 91325f8.

📒 Files selected for processing (11)

• cpp/tensorrt_llm/thop/moeOp.cpp (6 hunks)
• tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py (3 hunks)
• tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py (1 hunks)
• tensorrt_llm/_torch/models/modeling_nemotron_h.py (7 hunks)
• tensorrt_llm/_torch/modules/fused_moe/create_moe.py (4 hunks)
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (3 hunks)
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py (4 hunks)
• tensorrt_llm/_torch/modules/fused_moe/interface.py (2 hunks)
• tensorrt_llm/_torch/modules/fused_moe/quantization.py (7 hunks)
• tensorrt_llm/_torch/utils.py (3 hunks)
• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (7 hunks)

🧰 Additional context used

🧠 Learnings (18)

📓 Common learnings

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

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.

📚 Learning: 2025-10-20T16:54:09.824Z

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:6-6
Timestamp: 2025-10-20T16:54:09.824Z
Learning: In tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py, the import `from ...modules.mamba.layernorm_gated import _layer_norm_fwd` is correct and should not be changed to modules.fla.layernorm_gated. The _layer_norm_fwd function exists in both modules/mamba/layernorm_gated.py and modules/fla/layernorm_gated.py, but the mamba version is the intended implementation for this use case.

Applied to files:

• tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py
• tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.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/models/checkpoints/hf/nemotron_h_weight_mapper.py
• tensorrt_llm/_torch/models/modeling_nemotron_h.py
• tensorrt_llm/_torch/modules/fused_moe/interface.py

📚 Learning: 2025-08-14T23:23:27.449Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py
• tensorrt_llm/_torch/modules/fused_moe/quantization.py
• tensorrt_llm/_torch/models/modeling_nemotron_h.py
• cpp/tensorrt_llm/thop/moeOp.cpp

📚 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/modules/fused_moe/quantization.py

📚 Learning: 2025-08-09T20:57:04.084Z

Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/quantization.py

📚 Learning: 2025-09-29T15:14:28.503Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation with asserts for total size and TP divisibility.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/quantization.py
• cpp/tensorrt_llm/thop/moeOp.cpp

📚 Learning: 2025-09-29T15:14:28.503Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/quantization.py
• cpp/tensorrt_llm/thop/moeOp.cpp

📚 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/modules/fused_moe/quantization.py

📚 Learning: 2025-08-20T07:43:36.447Z

Learnt from: ChristinaZ
Repo: NVIDIA/TensorRT-LLM PR: 7068
File: cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh:169-172
Timestamp: 2025-08-20T07:43:36.447Z
Learning: In TensorRT-LLM MOE kernels, when processing up to 128 experts across 32 threads, each thread handles at most 4 experts (N < 5 constraint), where N represents candidates per thread rather than total system capacity.

Applied to files:

• cpp/tensorrt_llm/thop/moeOp.cpp

📚 Learning: 2025-08-17T15:07:01.420Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 6968
File: cpp/tensorrt_llm/thop/loraOp.cpp:133-141
Timestamp: 2025-08-17T15:07:01.420Z
Learning: In TensorRT-LLM's LoRA implementation, the LoraImpl::run() method handles setStream() internally in _runGemm(), along with setWorkspace(). Both stream and workspace are passed as arguments to run(), so there's no need to call setStream() explicitly in loraOp.cpp - this avoids redundancy and follows the intended architectural separation.

Applied to files:

• cpp/tensorrt_llm/thop/moeOp.cpp

📚 Learning: 2025-07-28T17:06:08.621Z

Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-08-29T14:07:45.863Z

Learnt from: EmmaQiaoCh
Repo: NVIDIA/TensorRT-LLM PR: 7370
File: tests/unittest/trt/model_api/test_model_quantization.py:24-27
Timestamp: 2025-08-29T14:07:45.863Z
Learning: In TensorRT-LLM's CI infrastructure, pytest skip markers (pytest.mark.skip) are properly honored even when test files have __main__ blocks that call test functions directly. The testing system correctly skips tests without requiring modifications to the __main__ block execution pattern.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-08-26T09:37:10.463Z

Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-09-09T09:40:45.658Z

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-08-26T09:37:10.463Z

Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-08-06T13:58:07.506Z

Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

📚 Learning: 2025-08-26T09:49:04.956Z

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Applied to files:

• tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

🧬 Code graph analysis (10)

tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py (1)

tensorrt_llm/_torch/utils.py (1)

• split (377-385)

tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py (2)

tensorrt_llm/_torch/utils.py (2)

• split (377-385)
• shape (139-140)

tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py (1)

• config (156-159)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py (2)

tensorrt_llm/_torch/utils.py (3)

• ActivationType (37-46)
• is_gated_activation (51-54)
• relu2 (388-389)

tensorrt_llm/_torch/modules/gated_mlp.py (1)

• GatedMLP (19-182)

tensorrt_llm/_torch/utils.py (3)

cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h (1)

• ActivationType (24-37)

tensorrt_llm/_torch/models/modeling_deepseekv3.py (1)

• split (157-162)

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_trtllm_moe.py (1)

• relu2 (85-86)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (2)

tensorrt_llm/_torch/utils.py (1)

• ActivationType (37-46)

cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h (1)

• ActivationType (24-37)

tensorrt_llm/_torch/modules/fused_moe/create_moe.py (1)

tensorrt_llm/_torch/utils.py (1)

• ActivationType (37-46)

tensorrt_llm/_torch/models/modeling_nemotron_h.py (8)

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py (1)

• Mamba2Metadata (88-137)

tensorrt_llm/_torch/utils.py (4)

• ActivationType (37-46)
• relu2 (388-389)
• shape (139-140)
• _ (226-232)

tensorrt_llm/_torch/attention_backend/interface.py (1)

• AttentionMetadata (44-396)

tensorrt_llm/_torch/modules/fused_moe/interface.py (3)

• MoEWeightLoadingMode (17-23)
• forward (528-570)
• _ (85-111)

tensorrt_llm/_torch/modules/fused_moe/create_moe.py (1)

• create_moe (60-216)

tensorrt_llm/_torch/modules/linear.py (1)

• Linear (1880-2105)

tensorrt_llm/_torch/modules/multi_stream_utils.py (1)

• maybe_execute_in_parallel (35-74)

tensorrt_llm/_torch/models/modeling_deepseekv3.py (1)

• DeepseekV3Gate (725-791)

cpp/tensorrt_llm/thop/moeOp.cpp (3)

cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp (2)

• getQuantParams (624-716)
• getQuantParams (624-626)

cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu (2)

• hidden_size (470-481)
• hidden_size (470-470)

cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_gemm_kernels.h (1)

• isGatedActivation (240-244)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

tensorrt_llm/_torch/utils.py (1)

• is_gated_activation (51-54)

tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (2)

tests/unittest/utils/util.py (1)

• skip_gpu_memory_less_than (200-206)

tests/scripts/perf-sanity/run_benchmark_serve.py (1)

• llm_models_root (174-175)

🪛 Ruff (0.14.5)

tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py

130-130: Avoid specifying long messages outside the exception class

(TRY003)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py

55-57: Avoid specifying long messages outside the exception class

(TRY003)

---

59-61: Avoid specifying long messages outside the exception class

(TRY003)

tensorrt_llm/_torch/models/modeling_nemotron_h.py

226-226: Unused method argument: kwargs

(ARG002)

---

427-427: Avoid specifying long messages outside the exception class

(TRY003)

tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

204-204: Avoid specifying long messages outside the exception class

(TRY003)

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

tensorrt_llm/_torch/utils.py (1)

8-8: Activation utilities, split, and relu2 are consistent with C++ and prior usage

• ActivationType plus is_gated_activation correctly mirror the C++ isGatedActivation (Swiglu, SwigluBias, Geglu), which is important for keeping Python/C++ MoE behavior in sync.
• Centralizing split here and using torch.split with a divisibility assert matches how it's used in the weight mappers and avoids per-model copies.
• relu2 matches the test helper implementation and is a good single source for the ReLU-squared activation.

No issues from my side.

Also applies to: 35-55, 377-389

tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py (1)

10-10: Using shared split utility is appropriate

Switching to tensorrt_llm._torch.utils.split keeps the Qwen3Next weight mapper aligned with the shared implementation and avoids duplicated helpers. Call sites match the new function signature and expected semantics.

Also applies to: 55-78, 89-101

cpp/tensorrt_llm/thop/moeOp.cpp (1)

863-899: Activation-aware expand_ratio in getQuantParams looks correct

The new base_activation_type parameter and expand_ratio = isGatedActivation(base_activation_type) ? 2 : 1 usage in the MXFP4 and NVFP4 quant branches correctly generalize the previous hard-coded factor of 2:

• For gated activations (Swiglu/SwigluBias/Geglu), you still enforce the 2× inter-size layout for fc1 weight scales.
• For non-gated activations (e.g., Relu2), the checks now expect the ungated inter size, matching the new activation-type support.

This is aligned with the plugin-side isGatedActivation behavior and the way inter_size is computed in the callers.

Also applies to: 900-959, 960-980, 1015-1080

tensorrt_llm/_torch/models/modeling_nemotron_h.py (3)

48-70: relu2-based MLPLayer wiring looks good

The MLPLayer now derives intermediate_size per-layer (handling both scalar and list configs) and uses relu2 as the activation, matching the new activation utilities. This is consistent with the broader ActivationType/R elu2 support and doesn’t change the DecoderLayer contract.

---

269-317: Aux stream dict plumbing into NemotronH layers is reasonable

• NemotronHModel creates three CUDA streams and maps them into self.aux_stream_dict for MoEShared, MoeChunkingOverlap, and MoeBalancer.
• NemotronHLayer now takes aux_stream_dict and passes it to NemotronHMOE when layer_type == "E", leaving existing Mamba/MLP/Transformer paths unchanged.

This keeps the multi-stream wiring localized and doesn’t affect non-MoE layers. Looks good.

Also applies to: 336-367

---

417-429: rms_norm_eps normalization is a good defensive addition

Normalizing rms_norm_eps from either rms_norm_eps or layer_norm_epsilon in the pretrained config before constructing the model ensures all downstream code can rely on config.rms_norm_eps being set. The fallback ValueError if neither is present is appropriate.

[Wanli-Jiang]

/bot run

[ZhanruiSunCh]

/bot run

[tensorrt-cicd]

PR_Github #25179 [ run ] triggered by Bot. Commit: 94e7884

[tensorrt-cicd]

PR_Github #25179 [ run ] completed with state SUCCESS. Commit: 94e7884
/LLM/main/L0_MergeRequest_PR pipeline #19038 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26039 [ run ] triggered by Bot. Commit: 5d91bd3

[tensorrt-cicd]

PR_Github #26039 [ run ] completed with state SUCCESS. Commit: 5d91bd3
/LLM/main/L0_MergeRequest_PR pipeline #19763 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26071 [ run ] triggered by Bot. Commit: 5d91bd3

[tensorrt-cicd]

PR_Github #26071 [ run ] completed with state FAILURE. Commit: 5d91bd3
/LLM/main/L0_MergeRequest_PR pipeline #19795 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26118 [ run ] triggered by Bot. Commit: 5d91bd3

[tensorrt-cicd]

PR_Github #26118 [ run ] completed with state SUCCESS. Commit: 5d91bd3
/LLM/main/L0_MergeRequest_PR pipeline #19834 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26198 [ run ] triggered by Bot. Commit: 5d91bd3

[tensorrt-cicd]

PR_Github #26198 [ run ] completed with state FAILURE. Commit: 5d91bd3
/LLM/main/L0_MergeRequest_PR pipeline #19901 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26238 [ run ] triggered by Bot. Commit: 5d91bd3

[tensorrt-cicd]

PR_Github #26238 [ run ] completed with state FAILURE. Commit: 5d91bd3
/LLM/main/L0_MergeRequest_PR pipeline #19935 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26305 [ run ] triggered by Bot. Commit: 5bfdd5b

[tensorrt-cicd]

PR_Github #26305 [ run ] completed with state DISABLED
L0 testing is limited to prioritized users. User Wanli-Jiang is not in the prioritized list. L0 testing cannot be triggered.

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26371 [ run ] triggered by Bot. Commit: 67ed6a4

[tensorrt-cicd]

PR_Github #26371 [ run ] completed with state SUCCESS. Commit: 67ed6a4
/LLM/main/L0_MergeRequest_PR pipeline #20030 completed with status: 'FAILURE'

[Wanli-Jiang]

/bot run --disable-fail-fast

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #26493 [ run ] triggered by Bot. Commit: 67ed6a4

[tensorrt-cicd]

PR_Github #26493 [ run ] completed with state SUCCESS. Commit: 67ed6a4
/LLM/main/L0_MergeRequest_PR pipeline #20143 completed with status: 'SUCCESS'