diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 8ca56a219dc3..67142b6e02d2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index ebab505ccb5c..e9dac388d34b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/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,16 +339,12 @@ 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 @@ -350,10 +352,7 @@ def _check_flashinfer_backend_support(self) -> bool: 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." diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py index 503106033ea6..775590ec2434 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/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: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a56816ea4832..ae074776870e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/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, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index be5a78639f9f..7fd0c65a6a8d 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/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] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 5ac2ad5f115f..b55138ab9608 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/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" diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 515dd979da6d..566349ee70cc 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/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] diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 91c37da09ea4..948c05cf162d 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/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] diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 9243fe82c041..bdccf5395cea 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/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)