Feature: Enable Flashinfer non-gated MoE bf16 (#43853)

Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com>

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),
         )

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()

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,