Enable trtllm-gen BF16 MoE for MTP (#24260)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -1020,9 +1020,11 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
                 "Please check flashinfer version to use bf16 with flashinfer_trtllm backend."
             ) from e
 
-    assert (
-        runner_config.activation == "silu"
-    ), "Only silu is supported for flashinfer trtllm moe"
+    _SUPPORTED_BF16_ACTIVATIONS = {"silu", "relu2"}
+    assert runner_config.activation in _SUPPORTED_BF16_ACTIVATIONS, (
+        f"Only {_SUPPORTED_BF16_ACTIVATIONS} are supported for flashinfer trtllm bf16 moe, "
+        f"got '{runner_config.activation}'."
+    )
     if not use_routed_topk:
         assert (
             dispatch_output.topk_output.topk_config.renormalize
@@ -1030,9 +1032,7 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
     assert (
         runner_config.num_fused_shared_experts == 0
     ), "Fused shared experts are not supported for flashinfer trtllm moe"
-    assert (
-        runner_config.is_gated
-    ), "Only gated MoEs are supported for flashinfer trtllm moe"
+    activation_type = get_activation_type(runner_config.activation)
 
     hidden_states = dispatch_output.hidden_states
     topk_output = dispatch_output.topk_output
@@ -1072,6 +1072,7 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
                     else 1.0
                 ),
                 tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]),
+                activation_type=activation_type,
             )
         else:
             assert TopKOutputChecker.format_is_bypassed(topk_output)
@@ -1094,6 +1095,7 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
                 routing_method_type=runner_config.routing_method_type,
                 routed_scaling_factor=runner_config.routed_scaling_factor,
                 tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]),
+                activation_type=activation_type,
             )
 
     return StandardCombineInput(hidden_states=final_hidden_states)

python/sglang/srt/server_args.py

@@ -1894,6 +1894,7 @@ class ServerArgs:
                         self.quantization
                         in ["fp8", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed"]
                         or is_kimi_k2_k25_thinking_int4
+                        or self.quantization is None
                     )
                 ):
                     self.moe_runner_backend = "flashinfer_trtllm"
@@ -3333,22 +3334,8 @@ class ServerArgs:
         ):
             self.speculative_draft_model_revision = "main"
 
-        # FlashInfer trtllm moe bf16 only support RenormalizeNaive routing method and Deepseek routing method
-        # It is hard to tell the routing method in draft model, and the moe layer in draft model is not the bottleneck among
-        # end to end, so we just avoid using trtllm_moe for speculative decoding.
-        from sglang.srt.layers.moe.utils import MoeRunnerBackend
-
         if self.speculative_moe_runner_backend is None:
-            self.speculative_moe_runner_backend = (
-                "auto"
-                if self.moe_runner_backend
-                in ["flashinfer_trtllm", "flashinfer_trtllm_routed"]
-                else self.moe_runner_backend
-            )
-        else:
-            assert not MoeRunnerBackend(
-                self.speculative_moe_runner_backend
-            ).is_flashinfer_trtllm(), "Currently speculative MoE runner backend doesn't support flashinfer_trtllm, please use triton or auto backend for speculative moe runner instead."
+            self.speculative_moe_runner_backend = self.moe_runner_backend
 
         if self.speculative_algorithm is not None:
             self.speculative_algorithm = self.speculative_algorithm.upper()