[FMDL-1328][feat] Add support for nano-v3 and super-v3 with pytorch backend #9261
+492
−142
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Wanli-Jiang
force-pushed
the
user/williamj/add-super-v3-pyt
branch
2 times, most recently
from
be8d8f6 to
e4e42e0
Compare
…ackend Signed-off-by: Wanli Jiang <[email protected]>
Wanli-Jiang
force-pushed
the
user/williamj/add-super-v3-pyt
branch
from
e4e42e0 to
91325f8
Compare
Contributor
📝 WalkthroughWalkthroughThe 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
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
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
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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.
Example instruction:
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
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
runMoeMinLantencykeeps 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 onisGatedActivation(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_typenear 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 tweaksThe 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_scaleand main-weight branches you rely onif weights[name].shape:to distinguish NVFP4 (tensor) vs FP8 (scalar). Making this explicit (e.g., checkingweights[name].ndim == 0) would be clearer and less brittle.- For the branches that slice
weights[name].shape[0] // 2, anassert weights[name].shape[0] % 2 == 0would 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 formoe_intermediate_sizelistsNemotronHMOE sets:
self.moe_intermediate_size = config.moe_intermediate_size[0] \ if isinstance(config.moe_intermediate_size, list) else config.moe_intermediate_sizeWhile MLPLayer treats list-valued
intermediate_sizeas:
- Use
intermediate_size[0]whenlen == 1(shared across layers).- Otherwise index by
layer_idx.If
config.moe_intermediate_sizecan 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_sizeIf 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 consistentA few points on the new NemotronHMOE:
Using
ActivationType.Relu2and passing it intocreate_moe(... activation_type=...)aligns with the new activation-type-aware MoE kernels.Gate (
DeepseekV3Gate) and experts (create_moewithMoEWeightLoadingMode.VANILLA) are wired the same way as other DeepSeekV3-style MoE modules, and the latent projection path is correctly guarded byuse_latent_moe.Pulling
all_rank_num_tokensfromattn_metadataand forwarding it intoself.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
📒 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.pytensorrt_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.pytensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_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.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/models/modeling_nemotron_h.pycpp/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.pycpp/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.pycpp/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
ActivationTypeplusis_gated_activationcorrectly mirror the C++isGatedActivation(Swiglu, SwigluBias, Geglu), which is important for keeping Python/C++ MoE behavior in sync.- Centralizing
splithere and usingtorch.splitwith a divisibility assert matches how it's used in the weight mappers and avoids per-model copies.relu2matches 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 sharedsplitutility is appropriateSwitching to
tensorrt_llm._torch.utils.splitkeeps 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 correctThe new
base_activation_typeparameter andexpand_ratio = isGatedActivation(base_activation_type) ? 2 : 1usage 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
isGatedActivationbehavior and the wayinter_sizeis 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 goodThe MLPLayer now derives
intermediate_sizeper-layer (handling both scalar and list configs) and usesrelu2as 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_dictfor MoEShared, MoeChunkingOverlap, and MoeBalancer.- NemotronHLayer now takes
aux_stream_dictand passes it to NemotronHMOE whenlayer_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 additionNormalizing
rms_norm_epsfrom eitherrms_norm_epsorlayer_norm_epsilonin the pretrained config before constructing the model ensures all downstream code can rely onconfig.rms_norm_epsbeing set. The fallbackValueErrorif neither is present is appropriate.
Comment on lines
+128
to
147
| if self.activation_type == ActivationType.Relu2: | ||
| module_list[expert_idx] = MLP( | ||
| hidden_size=self.hidden_size, | ||
| intermediate_size=self.intermediate_size, | ||
| bias=False, | ||
| activation=relu2, | ||
| dtype=self.dtype, | ||
| config=model_config, | ||
| layer_idx=None, | ||
| ) | ||
| else: | ||
| module_list[expert_idx] = GatedMLP( | ||
| hidden_size=self.hidden_size, | ||
| intermediate_size=self.intermediate_size, | ||
| bias=False, | ||
| dtype=self.dtype, | ||
| config=model_config, | ||
| reduce_output=False, | ||
| ) | ||
| else: |
Contributor
There was a problem hiding this comment.
Relu2 experts break weight loading
When activation_type == ActivationType.Relu2 we now instantiate MLP, but load_weights() still unconditionally reaches into each expert’s gate_up_proj/down_proj. MLP doesn’t expose gate_up_proj, so as soon as weights are loaded this will raise an AttributeError. Please keep experts in the Relu2 branch compatible with the existing weight-loading path (e.g., branch load_weights/pack_params for the non-gated case or use a module that still surfaces gate_up_proj).
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py around lines
128-147, the Relu2 branch instantiates MLP which does not expose
gate_up_proj/gate_down_proj expected by the existing load_weights/pack_params
path and thus will raise AttributeError; fix by making the Relu2 experts
compatible with the loader: either (A) instantiate or wrap a module that
provides gate_up_proj and gate_down_proj attributes (pointing to the MLP's
up_proj/down_proj or to dummy aliases) so the existing weight-loading code can
access them unchanged, or (B) modify load_weights/pack_params to detect the
non-gated MLP (use hasattr checks) and follow the non-gated parameter path;
implement one of these two approaches consistently so weight loading no longer
assumes gate_* attributes for Relu2 experts.
Comment on lines
+66
to
+71
| if model_folder == "Nemotron-H-8B-Base-8K": | ||
| skip_gpu_memory_less_than( | ||
| (2 * 8 + 1) * 2**30)() # 8B, bf16, plus 1 GB for good measure | ||
| elif model_folder == "Nemotron-Nano-3-30B-A3.5B-dev-1024": | ||
| skip_gpu_memory_less_than( | ||
| (2 * 30 + 1) * 2**30)() # 30B, bf16, plus 1 GB for good measure |
Contributor
There was a problem hiding this comment.
Fix runtime skip handling
skip_gpu_memory_less_than returns a pytest.mark.skipif decorator. Invoking it as ...(), without passing a function, raises TypeError: __call__() missing 1 required positional argument: 'func'. On hosts that don’t meet the memory requirement we’ll blow up instead of skipping the test. Please evaluate the condition directly and call pytest.skip(...) (or apply the mark during parametrization) so the test actually skips.
@@
- if model_folder == "Nemotron-H-8B-Base-8K":
- skip_gpu_memory_less_than(
- (2 * 8 + 1) * 2**30)() # 8B, bf16, plus 1 GB for good measure
- elif model_folder == "Nemotron-Nano-3-30B-A3.5B-dev-1024":
- skip_gpu_memory_less_than(
- (2 * 30 + 1) * 2**30)() # 30B, bf16, plus 1 GB for good measure
+ if model_folder == "Nemotron-H-8B-Base-8K":
+ required_bytes = (2 * 8 + 1) * 2**30 # 8B, bf16, plus 1 GB for good measure
+ elif model_folder == "Nemotron-Nano-3-30B-A3.5B-dev-1024":
+ required_bytes = (2 * 30 + 1) * 2**30 # 30B, bf16, plus 1 GB for good measure
+ else:
+ required_bytes = None
+
+ if required_bytes is not None and required_bytes > torch.cuda.get_device_properties(0).total_memory:
+ pytest.skip(
+ f"Not enough GPU memory for this test (wanted {required_bytes}, "
+ f"have {torch.cuda.get_device_properties(0).total_memory})")🤖 Prompt for AI Agents
In tests/unittest/_torch/modeling/test_modeling_nemotron_h.py around lines 66 to
71, the code is calling skip_gpu_memory_less_than(... )() which incorrectly
invokes the pytest.mark.skipif decorator and triggers a TypeError; instead
evaluate the same memory condition directly and call pytest.skip("insufficient
GPU memory for <model>") when the condition is true (or apply the skip mark
during parametrization). Replace the current skip_gpu_memory_less_than(... )()
calls with a runtime check of the required GPU memory and a pytest.skip call
with a clear message (or attach the returned mark to the test function/param) so
the test is properly skipped rather than raising.
yuxianq
reviewed
Nov 19, 2025
| TORCH_CHECK(fc2_weight_block.dim() == 3, "fc2 weight block must be 3D"); | ||
| TORCH_CHECK(fc2_global.dim() == 1, "fc2 global must be 1D"); | ||
| // Check shapes | ||
| int expand_ratio = isGatedActivation(base_activation_type) ? 2 : 1; |
Collaborator
There was a problem hiding this comment.
Van we extract this line to outer scope instead of repeating it for each branch?
yuxianq
reviewed
Nov 19, 2025
| activation=relu2, | ||
| dtype=self.dtype, | ||
| config=model_config, | ||
| layer_idx=None, |
Collaborator
There was a problem hiding this comment.
Shoule we use layer_idx=self.layer_idx instead? self.layer_idx seems unused now.
yuxianq
reviewed
Nov 19, 2025
| bias=self.mlp_bias, | ||
| activation_type=self.activation_type, | ||
| # Default values | ||
| override_quant_config=None, |
Collaborator
There was a problem hiding this comment.
Can we remove line L191~195 since they just use default values?
yuxianq
reviewed
Nov 19, 2025
| assert completions_batching[i] == expected_completions[i] | ||
| assert completions_no_batching[i] == expected_completions[i] | ||
| else: | ||
| assert similar(completions_batching[i], |
Collaborator
There was a problem hiding this comment.
Why not compare expected completion in this case?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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
Improvements
Tests
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.
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 thestage-listparameter 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.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.