diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 734883526..e2b0f08c4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8db2115e6..fc1c961da 100644 --- a/python/sglang/srt/server_args.py +++ b/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()