[Kernel][MoE] Allow FlashInfer MXINT4 MoE for gated SiLU (#46518)
Signed-off-by: Kaihang Jiang <kaihangj@nvidia.com>
Kaihang Jiang committed
Jun 24, 2026 at 19:32 UTC
fc7fc421e98863c4ffb1aa02d46bd6e4d0202c26
2 files changed
+16
-2
tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
+12
@@ -5,9 +5,13 @@
5
import pytest
6
import torch
7
8
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
10
fused_marlin_moe,
11
)
12
+from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
13
+ TrtLlmMxint4ExpertsMonolithic,
14
+)
15
from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
16
grouped_topk,
17
)
@@ -77,6 +81,14 @@ __all__ = [
81
]
82
83
84
+def test_trtllm_mxint4_activation_supports_vllm_gated_silu():
85
+ assert TrtLlmMxint4ExpertsMonolithic._supports_activation(MoEActivation.SILU)
86
+ assert TrtLlmMxint4ExpertsMonolithic._supports_activation(MoEActivation.SWIGLUOAI)
87
+ assert not TrtLlmMxint4ExpertsMonolithic._supports_activation(
88
+ MoEActivation.RELU2_NO_MUL
89
+ )
90
+
91
+
92
def marlin_quantize_moe_weights(
93
weights_bf16: torch.Tensor, group_size: int = 32
94
) -> tuple[torch.Tensor, torch.Tensor]:
vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
+4
-2
@@ -68,8 +68,10 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
68
69
@staticmethod
70
def _supports_activation(activation: MoEActivation) -> bool:
71
- # FlashInfer MxInt4 uses a fused SwiGLU activation.
72
- return activation == MoEActivation.SWIGLUOAI
71
+ # FlashInfer MxInt4 names the standard gated SiLU path "SwiGLU".
72
+ # In vLLM MoE configs that maps to SILU/silu_and_mul; SWIGLUOAI is
73
+ # kept as an alias for consistency with other FlashInfer backends.
74
+ return activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI)
75
76
@staticmethod
77
def _supports_parallel_config(