[TRTLLM-13052][feat] Enable TRTLLM moe backend for nemotron-h BF16 ckpt

tensorrt_llm/_torch/modules/fused_moe/create_moe.py

@@ -183,9 +183,7 @@ def resolve_moe_cls(
     has_quant = (effective_quant_config is not None
                  and effective_quant_config.layer_quant_mode.has_any_quant(
                      exclude_kv_cache=True))
-    if (moe_cls == TRTLLMGenFusedMoE and not has_quant
-            and not TRTLLMGenFusedMoE._supports_flashinfer_bf16_routing_method(
-                routing_method)):
+    if (moe_cls == TRTLLMGenFusedMoE and not has_quant):
         moe_cls = CutlassFusedMoE
 
     # Routed-expert LoRA is supported only on CutlassFusedMoE with unquantized

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -107,6 +107,12 @@ class TRTLLMGenFusedMoE(MoE):
         QuantAlgo.W4A8_MXFP4_MXFP8,
     }
 
+    # Activations supported by the FlashInfer BF16 kernels: Swiglu and Relu2.
+    _BF16_SUPPORTED_ACTIVATIONS = {
+        ActivationType.Swiglu,
+        ActivationType.Relu2,
+    }
+
     @classmethod
     def can_implement(
         cls,
@@ -333,27 +339,20 @@ def _is_unquantized_path(self) -> bool:
         return self.quant_config is None or not self.quant_config.layer_quant_mode.has_any_quant(
             exclude_kv_cache=True)
 
-    @staticmethod
-    def _supports_flashinfer_bf16_routing_method(
-        routing_method: BaseMoeRoutingMethod, ) -> bool:
-        # FIXME: ban DeepSeekV3 FlashInfer trtllm_bf16_routed_moe() as it appears to have bug
-        return not isinstance(routing_method, DeepSeekV3MoeRoutingMethod)
-
     def _requires_separated_routing(self) -> bool:
-        """Whether this backend instance expects precomputed top-k routing."""
-        # FIXME: ban FlashInfer BF16 MoE direct routing as it appears to have accuracy bug
-        return self.use_flashinfer and self._is_unquantized_path()
+        """BF16 FlashInfer uses separated routing, except DeepSeekV3 which uses
+        the fused kernel (its separated variant has accuracy issues)."""
+        if not (self.use_flashinfer and self._is_unquantized_path()):
+            return False
+        return not isinstance(self.routing_method, DeepSeekV3MoeRoutingMethod)
 
     def _check_flashinfer_backend_support(self) -> bool:
         # For BF16 (unquantized) path, we will use FlashInfer regardless whether
         # env TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER=1 is set or not as it's the only way.
         if self._is_unquantized_path():
             if not self._is_flashinfer_fused_moe_available():
                 return False
-            if self.activation_type != ActivationType.Swiglu:
-                return False
-            if not self._supports_flashinfer_bf16_routing_method(
-                    self.routing_method):
+            if self.activation_type not in self._BF16_SUPPORTED_ACTIVATIONS:
                 return False
             return True
 
@@ -451,8 +450,10 @@ def _check_configs(self):
             "TRTLLMGenFusedMoE only supports bf16 (FlashInfer), fp8_block_scaling, nvfp4, w4a16_mxfp4, w4a8_mxfp4_fp8 and w4a8_mxfp4_mxfp8 dtypes."
 
         if not self.has_any_quant:
-            assert self.activation_type == ActivationType.Swiglu, \
-                "TRTLLMGenFusedMoE BF16 path only supports Swiglu activation."
+            assert self.activation_type in self._BF16_SUPPORTED_ACTIVATIONS, \
+                ("TRTLLMGenFusedMoE BF16 path only supports "
+                 f"{[a.name for a in self._BF16_SUPPORTED_ACTIVATIONS]} activations, "
+                 f"got {self.activation_type.name}.")
             assert not self.bias and self.swiglu_alpha is None and self.swiglu_beta is None and self.swiglu_limit is None, \
                 "TRTLLMGenFusedMoE BF16 path does not support bias/swiglu custom parameters."
 

tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py

@@ -25,8 +25,6 @@
 
 import torch
 
-from ...utils import ActType_TrtllmGen
-
 # Global registry for MoE backends
 _MOE_OP_BACKEND_REGISTRY: Dict[str, Type["MoEOpBackend"]] = {}
 
@@ -805,11 +803,8 @@ def run_bf16_moe(
         enable_pdl=None,
         tune_max_num_tokens=8192,
     ):
-        # FlashInfer BF16 MoE does not expose an activation_type argument.
-        # TRTLLMGen constrains the BF16 path to Swiglu, so reject anything
-        # else here instead of silently calling a mismatched kernel.
-        if gated_act_type != ActType_TrtllmGen.SwiGlu:
-            raise ValueError("FlashInfer BF16 fused MoE only supports Swiglu activation.")
+        # Forward the activation (Swiglu/Relu2) to the FlashInfer BF16 kernels.
+        activation_type = self.cvt_activation_type(gated_act_type)
 
         if router_logits is not None:
             result = self._fused_moe.trtllm_bf16_moe(
@@ -832,6 +827,7 @@ def run_bf16_moe(
                 do_finalize=do_finalize,
                 enable_pdl=enable_pdl,
                 tune_max_num_tokens=tune_max_num_tokens,
+                activation_type=activation_type,
             )
         else:
             packed_topk_ids = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
@@ -856,6 +852,7 @@ def run_bf16_moe(
                 do_finalize=do_finalize,
                 enable_pdl=enable_pdl,
                 tune_max_num_tokens=tune_max_num_tokens,
+                activation_type=activation_type,
             )
 
         if output is not None and do_finalize:

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -6816,6 +6816,35 @@ def test_auto_dtype_4gpus(self, tp_size, ep_size, attention_dp,
             task.evaluate(llm,
                           extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
 
+    @skip_pre_blackwell
+    @pytest.mark.skip_less_device_memory(80000)
+    @pytest.mark.skip_less_mpi_world_size(4)
+    @parametrize_with_ids("attention_dp", [False, True])
+    def test_bf16_trtllm_gen_moe_backend(self, attention_dp):
+
+        kv_cache_config = KvCacheConfig(enable_block_reuse=False,
+                                        mamba_ssm_cache_dtype="float32")
+        pytorch_config = dict(disable_overlap_scheduler=False,
+                              cuda_graph_config=CudaGraphConfig(
+                                  max_batch_size=32, enable_padding=True))
+
+        with LLM(
+                f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-BF16",
+                kv_cache_config=kv_cache_config,
+                max_batch_size=32,
+                tensor_parallel_size=4,
+                moe_expert_parallel_size=4,
+                enable_attention_dp=attention_dp,
+                moe_config=MoeConfig(backend="TRTLLM"),
+                **pytorch_config,
+        ) as llm:
+            task = MMLU(self.MODEL_NAME)
+            task.evaluate(llm,
+                          extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
+            task = GSM8K(self.MODEL_NAME)
+            task.evaluate(llm,
+                          extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
+
     def _run_nvfp4_4gpus_eplb(self, moe_backend, eplb_config, model_path):
         kv_cache_config = KvCacheConfig(
             enable_block_reuse=False,

tests/integration/test_lists/qa/llm_function_core.txt

@@ -699,6 +699,8 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP8_PP1]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP8_PP1_ADP]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP_MTP]
+accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_bf16_trtllm_gen_moe_backend[attention_dp=True]
+accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_bf16_trtllm_gen_moe_backend[attention_dp=False]
 accuracy/test_llm_api_pytorch.py::TestPhi4MiniInstruct::test_auto_dtype
 accuracy/test_llm_api_pytorch.py::TestQwen2_7BInstruct::test_auto_dtype
 accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[latency]

tests/integration/test_lists/test-db/l0_b200.yml

@@ -114,6 +114,7 @@ l0_b200:
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL"
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM"
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DENSEGEMM"
+  - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe
   # ------------- MoE: test_single_gpu (by backend) ---------------
   - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS"
   - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM"

tests/integration/test_lists/test-db/l0_b300.yml

@@ -40,6 +40,7 @@ l0_b300:
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM"
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL"
   - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM"
+  - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe
   # ------------- MoE: test_single_gpu (specific quant per backend) ---------------
   # CUTLASS backend: FP8, NVFP4, W4A8_MXFP4_MXFP8, W8A16
   - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize]

tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

@@ -46,6 +46,8 @@ l0_gb200_multi_gpus:
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP]
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP_MTP]
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=CUTEDSL]
+  - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_bf16_trtllm_gen_moe_backend[attention_dp=True]
+  - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_bf16_trtllm_gen_moe_backend[attention_dp=False]
   - accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_tp4[torch_compile=False]
   - accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_tp4[torch_compile=True]
   - accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=False]

tests/unittest/_torch/modules/moe/test_moe_backend.py

@@ -51,7 +51,10 @@
 
 from tensorrt_llm._torch.autotuner import AutoTuner, autotune
 from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod
+from tensorrt_llm._torch.modules.fused_moe import (
+    DeepSeekV3MoeRoutingMethod,
+    RenormalizeMoeRoutingMethod,
+)
 from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend
 from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoEWeightLoadingMode
 from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm
@@ -744,3 +747,131 @@ def run_moe():
             with torch.inference_mode():
                 output = run_moe()
                 ref_fused_moe.check_accuracy(output, ref_output)
+
+
+# ============================================================================
+# BF16 (unquantized) TRTLLM-Gen MoE: DeepSeekV3 / Renormalize routing
+# ============================================================================
+# The main test_moe_backend skips TRTLLM + quant_algo=None, so cover the BF16
+# FlashInfer path here (Nemotron-H enablement): DeepSeekV3/Renormalize routing
+# x Relu2/Swiglu, via both fused and separated routing.
+
+# DeepSeekV3 trtllm-gen routing requires num_experts >= 22, multiple of 4.
+_BF16_UNQUANT_NUM_EXPERTS = 72
+_BF16_UNQUANT_TOP_K = 6
+_BF16_UNQUANT_HIDDEN = 1024
+_BF16_UNQUANT_INTERMEDIATE = 512
+
+
+def _make_bf16_routing_method(routing_kind: str, top_k: int, num_experts: int, device: str):
+    if routing_kind == "renormalize":
+        return RenormalizeMoeRoutingMethod(top_k=top_k)
+    # DeepSeekV3 (noaux_tc): sigmoid scores + correction bias, single group.
+    bias = torch.randn(num_experts, dtype=torch.float32, device=device)
+    return DeepSeekV3MoeRoutingMethod(
+        top_k=top_k,
+        n_group=1,
+        topk_group=1,
+        routed_scaling_factor=2.5,
+        callable_e_score_correction_bias=lambda: bias,
+    )
+
+
+@pytest.mark.parametrize(
+    "trtllm_use_router_logits", [True, False], ids=["fused_routing", "separated_routing"]
+)
+@pytest.mark.parametrize("seq_len", [8, 256])
+@pytest.mark.parametrize(
+    "activation_type", [ActivationType.Relu2, ActivationType.Swiglu], ids=["relu2", "swiglu"]
+)
+@pytest.mark.parametrize("routing_kind", ["deepseekv3", "renormalize"])
+def test_trtllm_bf16_unquantized_moe(
+    routing_kind, activation_type, seq_len, trtllm_use_router_logits
+):
+    """TRTLLM-Gen BF16 (unquantized) MoE accuracy vs the reference impl."""
+    backend_type = MoeBackendType.TRTLLM
+    dtype = torch.bfloat16
+
+    can_impl, skip_reason = get_backend_class(backend_type).can_implement(
+        None, dtype_activation=dtype
+    )
+    if not can_impl:
+        pytest.skip(skip_reason)
+
+    num_experts = _BF16_UNQUANT_NUM_EXPERTS
+    top_k = _BF16_UNQUANT_TOP_K
+    hidden_size = _BF16_UNQUANT_HIDDEN
+    intermediate_size = _BF16_UNQUANT_INTERMEDIATE
+
+    skip_if_insufficient_gpu_memory(num_experts, hidden_size, intermediate_size, dtype)
+
+    mapping = Mapping()
+    mapping.rank = mpi_rank()
+
+    with torch.device(f"cuda:{mapping.rank}"):
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        AutoTuner.get().setup_distributed_state(mapping)
+
+        routing_method = _make_bf16_routing_method(routing_kind, top_k, num_experts, "cuda")
+
+        x = torch.randn((seq_len, hidden_size), dtype=dtype, device="cuda")
+        router_logits = torch.randn((seq_len, num_experts), dtype=dtype, device="cuda")
+
+        # Unquantized path: get_test_quant_params returns BaseQuantizeUtil.
+        quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params(None, x, backend_type)
+        quantize_util = quantize_util_cls(
+            num_experts=num_experts,
+            dtype=dtype,
+            intermediate_size=intermediate_size,
+            hidden_size=hidden_size,
+            quant_config=quant_config,
+            activation_type=activation_type,
+        )
+        weights = quantize_util.create_weights(**quant_kwargs)
+
+        backend = create_test_backend(
+            backend_type=backend_type,
+            routing_method=routing_method,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            dtype=dtype,
+            quant_config=quant_config,
+            mapping=mapping,
+            activation_type=activation_type,
+        )
+        backend.load_weights([weights])
+        backend.post_load_weights()
+        backend.cuda()
+
+        ref_fused_moe = quantize_util.create_ref_module(routing_method)
+        ref_fused_moe.load_weights([weights])
+        ref_fused_moe.cuda()
+
+        with torch.inference_mode():
+            ref_output = ref_fused_moe.forward(x, router_logits)
+
+        AutoTuner.get().clear_cache()
+
+        def run_moe():
+            token_selected_experts, token_final_scales = routing_method.apply(router_logits)
+            x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False)
+            return run_backend_moe(
+                backend,
+                backend_type,
+                x_quantized,
+                x_sf,
+                token_selected_experts,
+                token_final_scales,
+                dtype,
+                router_logits=router_logits,
+                trtllm_use_router_logits=trtllm_use_router_logits,
+            )
+
+        # Autotune, then verify accuracy against the reference.
+        with torch.inference_mode(), autotune(cache_path="/tmp/moe_autotuner_cache.json"):
+            _ = run_moe()
+        with torch.inference_mode():
+            output = run_moe()
+            ref_fused_moe.check_accuracy(output, ref_output)