diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 43126195205..9ae82d4e73c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -317,9 +317,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout assert not apply_router_weight_on_input - assert activation == MoEActivation.SILU + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + activation_type = activation_to_flashinfer_int(activation) assert self.topk <= global_num_experts - assert self.topk <= 10 assert global_num_experts % 4 == 0 assert self.quant_config.block_shape in [[128, 128], [1, 32]] # Kernel expects #experts <= #threads 512 @@ -333,13 +333,19 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit use_shuffled_weight = True weight_layout = WeightLayout.MajorK hidden_states_scale = a1q_scale + # FlashInfer expects None for non-grouped MXFP8 routing configs. + n_group = num_expert_group or None + selected_topk_group = topk_group or None else: + assert self.topk <= 10 fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 use_shuffled_weight = True weight_layout = WeightLayout.BlockMajorK hidden_states_scale = a1q_scale.t().contiguous() + n_group = num_expert_group or 0 + selected_topk_group = topk_group or 0 - return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( + kwargs = dict( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -350,8 +356,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, top_k=self.topk, - n_group=(num_expert_group or 0), - topk_group=(topk_group or 0), + n_group=n_group, + topk_group=selected_topk_group, intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, @@ -361,6 +367,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, ) + if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL: + kwargs["activation_type"] = activation_type + return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs) def _apply_per_tensor( self,