A high-throughput and memory-efficient inference and serving engine for LLMs
Feature: Enable Flashinfer non-gated MoE bf16 (#43853)
Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com>
amirkl94 committed
Jun 17, 2026 at 17:32 UTC
8b2b566ea710e48149886aecf2e17a2c78d55c11
3 files changed
+30
-8
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+10
-3
@@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
RoutingMethodType,
13
)
14
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
15
+ activation_to_flashinfer_int,
16
+)
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
QuantKey,
19
)
@@ -54,8 +57,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
57
58
@staticmethod
59
def _supports_no_act_and_mul() -> bool:
57
- """BF16 kernels do not support non-gated MoE"""
58
- return False
60
+ """BF16 kernels support non-gated MoE via RELU2_NO_MUL."""
61
+ return True
62
63
@staticmethod
64
def _supports_quant_scheme(
@@ -67,7 +70,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
70
71
@staticmethod
72
def _supports_activation(activation: MoEActivation) -> bool:
70
- return activation in [MoEActivation.SILU]
73
+ """Supports SiLU (gated) and RELU^2 (non-gated) activations."""
74
+ return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
75
76
@staticmethod
77
def _supports_routing_method(
@@ -123,6 +127,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
127
) -> torch.Tensor:
128
import flashinfer
129
130
+ assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
131
+
132
return flashinfer.fused_moe.trtllm_bf16_moe(
133
routing_logits=router_logits,
134
routing_bias=e_score_correction_bias,
@@ -138,4 +144,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
144
local_num_experts=self.local_num_experts,
145
routed_scaling_factor=routed_scaling_factor,
146
routing_method_type=self.routing_method_type,
147
+ activation_type=activation_to_flashinfer_int(activation),
148
)
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+12
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (
19
FusedMoEQuantConfig,
20
)
21
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22
+ align_moe_weights_for_fi,
23
convert_moe_weights_to_flashinfer_trtllm_block_layout,
24
swap_w13_to_w31,
25
)
@@ -269,11 +270,22 @@ def convert_to_unquantized_kernel_format(
270
w13_weight = swap_w13_to_w31(w13_weight)
271
272
elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
273
+ is_act_and_mul = layer.moe_config.is_act_and_mul
274
+ if not is_act_and_mul:
275
+ # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK
276
+ # weight layout uses block_k=128). Pad along the intermediate dim when
277
+ # the model + TP split don't satisfy the constraint.
278
+ w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi(
279
+ w13_weight, w2_weight, is_act_and_mul, min_alignment=128
280
+ )
281
+ layer.moe_config.intermediate_size_per_partition = padded_intermediate
282
+
283
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
284
w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
285
_cache_permute_indices,
286
w13_weight,
287
w2_weight,
288
+ is_gated_act_gemm=is_act_and_mul,
289
)
290
291
return w13_weight.contiguous(), w2_weight.contiguous()
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+8
-5
@@ -108,6 +108,7 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
108
cache_permute_indices: dict[torch.Size, torch.Tensor],
109
w13_weight: torch.Tensor,
110
w2_weight: torch.Tensor,
111
+ is_gated_act_gemm: bool = True,
112
) -> tuple[torch.Tensor, torch.Tensor]:
113
"""Convert expert weights to FlashInfer's block layout.
114
@@ -166,9 +167,11 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
167
cache_permute_indices,
168
w13_expert_uint8,
169
epilogue_tile_m,
170
+ is_gated_act_gemm=is_gated_act_gemm,
171
)
170
- rows = w13_expert_uint8.shape[0]
171
- permute_indices = (permute_indices + rows // 2) % rows
172
+ if is_gated_act_gemm:
173
+ rows = w13_expert_uint8.shape[0]
174
+ permute_indices = (permute_indices + rows // 2) % rows
175
_copy_permuted_expert_to_block_layout(
176
w13_weights_shuffled_tensor[i],
177
w13_expert_uint8,
@@ -288,12 +291,12 @@ def align_trtllm_fp4_moe_hidden_dim_for_fi(
291
return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_hidden_size
292
293
291
-def align_fp8_moe_weights_for_fi(
294
+def align_moe_weights_for_fi(
295
w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16
296
) -> tuple[torch.Tensor, torch.Tensor, int]:
297
"""Pad intermediate size so FlashInfer kernels' alignment constraints hold.
298
296
- Some FlashInfer FP8 MoE kernels require the (gated) intermediate size
299
+ Some FlashInfer MoE kernels require the (gated) intermediate size
300
used for GEMM to be divisible by a small alignment value. When this is
301
not satisfied (e.g. with certain tensor-parallel sizes), we pad the
302
gate/up and down projection weights along the intermediate dim.
@@ -492,7 +495,7 @@ def prepare_fp8_moe_layer_for_fi(
495
# for the gate-up proj. Pad the weights to respect this.
496
if not block_quant:
497
min_alignment = 16 if is_gated else 128
495
- w13, w2, new_intermediate = align_fp8_moe_weights_for_fi(
498
+ w13, w2, new_intermediate = align_moe_weights_for_fi(
499
w13,
500
w2,
501
layer.moe_config.is_act_and_mul,