From 8b2b566ea710e48149886aecf2e17a2c78d55c11 Mon Sep 17 00:00:00 2001 From: amirkl94 <203507526+amirkl94@users.noreply.github.com> Date: Wed, 17 Jun 2026 17:32:49 +0300 Subject: [PATCH] Feature: Enable Flashinfer non-gated MoE bf16 (#43853) Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com> --- .../layers/fused_moe/experts/trtllm_bf16_moe.py | 13 ++++++++++--- .../layers/fused_moe/oracle/unquantized.py | 12 ++++++++++++ .../layers/quantization/utils/flashinfer_utils.py | 13 ++++++++----- 3 files changed, 30 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 033b7dc7d93..550d6b5341d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, ) @@ -54,8 +57,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): @staticmethod def _supports_no_act_and_mul() -> bool: - """BF16 kernels do not support non-gated MoE""" - return False + """BF16 kernels support non-gated MoE via RELU2_NO_MUL.""" + return True @staticmethod def _supports_quant_scheme( @@ -67,7 +70,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in [MoEActivation.SILU] + """Supports SiLU (gated) and RELU^2 (non-gated) activations.""" + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] @staticmethod def _supports_routing_method( @@ -123,6 +127,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): ) -> torch.Tensor: import flashinfer + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + return flashinfer.fused_moe.trtllm_bf16_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -138,4 +144,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): local_num_experts=self.local_num_experts, routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, + activation_type=activation_to_flashinfer_int(activation), ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 5c156dc9b16..a7dcd801376 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + align_moe_weights_for_fi, convert_moe_weights_to_flashinfer_trtllm_block_layout, swap_w13_to_w31, ) @@ -269,11 +270,22 @@ def convert_to_unquantized_kernel_format( w13_weight = swap_w13_to_w31(w13_weight) elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + is_act_and_mul = layer.moe_config.is_act_and_mul + if not is_act_and_mul: + # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK + # weight layout uses block_k=128). Pad along the intermediate dim when + # the model + TP split don't satisfy the constraint. + w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi( + w13_weight, w2_weight, is_act_and_mul, min_alignment=128 + ) + layer.moe_config.intermediate_size_per_partition = padded_intermediate + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout( _cache_permute_indices, w13_weight, w2_weight, + is_gated_act_gemm=is_act_and_mul, ) return w13_weight.contiguous(), w2_weight.contiguous() diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 9e29ddaeb6f..1cbfdf69c99 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -108,6 +108,7 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout( cache_permute_indices: dict[torch.Size, torch.Tensor], w13_weight: torch.Tensor, w2_weight: torch.Tensor, + is_gated_act_gemm: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """Convert expert weights to FlashInfer's block layout. @@ -166,9 +167,11 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout( cache_permute_indices, w13_expert_uint8, epilogue_tile_m, + is_gated_act_gemm=is_gated_act_gemm, ) - rows = w13_expert_uint8.shape[0] - permute_indices = (permute_indices + rows // 2) % rows + if is_gated_act_gemm: + rows = w13_expert_uint8.shape[0] + permute_indices = (permute_indices + rows // 2) % rows _copy_permuted_expert_to_block_layout( w13_weights_shuffled_tensor[i], w13_expert_uint8, @@ -288,12 +291,12 @@ def align_trtllm_fp4_moe_hidden_dim_for_fi( return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_hidden_size -def align_fp8_moe_weights_for_fi( +def align_moe_weights_for_fi( w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16 ) -> tuple[torch.Tensor, torch.Tensor, int]: """Pad intermediate size so FlashInfer kernels' alignment constraints hold. - Some FlashInfer FP8 MoE kernels require the (gated) intermediate size + Some FlashInfer MoE kernels require the (gated) intermediate size used for GEMM to be divisible by a small alignment value. When this is not satisfied (e.g. with certain tensor-parallel sizes), we pad the gate/up and down projection weights along the intermediate dim. @@ -492,7 +495,7 @@ def prepare_fp8_moe_layer_for_fi( # for the gate-up proj. Pad the weights to respect this. if not block_quant: min_alignment = 16 if is_gated else 128 - w13, w2, new_intermediate = align_fp8_moe_weights_for_fi( + w13, w2, new_intermediate = align_moe_weights_for_fi( w13, w2, layer.moe_config.is_act_and_mul,