[MoE] Plumb gemm1_alpha/beta/clamp_limit into TRT-LLM FP8 MoE (#45723)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Jul 1, 2026 at 14:34 UTC
fa248139a0206b3e39780296e3f056a978957f63
4 files changed
+52
-3
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+46
-3
@@ -56,6 +56,35 @@ class TrtLlmFp8ExpertsBase:
56
self.moe_config = moe_config
57
self.quant_config = quant_config
58
59
+ # Per-expert SwiGLU parameters from quant_config (MXFP8 + Swiglu only).
60
+ device = torch.accelerator.current_device_index()
61
+ if quant_config.gemm1_alpha is not None:
62
+ self.gemm1_alpha = torch.tensor(
63
+ [quant_config.gemm1_alpha] * self.local_num_experts,
64
+ dtype=torch.float32,
65
+ device=device,
66
+ )
67
+ else:
68
+ self.gemm1_alpha = None
69
+
70
+ if quant_config.gemm1_beta is not None:
71
+ self.gemm1_beta = torch.tensor(
72
+ [quant_config.gemm1_beta] * self.local_num_experts,
73
+ dtype=torch.float32,
74
+ device=device,
75
+ )
76
+ else:
77
+ self.gemm1_beta = None
78
+
79
+ if quant_config.gemm1_clamp_limit is not None:
80
+ self.gemm1_clamp_limit = torch.tensor(
81
+ [quant_config.gemm1_clamp_limit] * self.local_num_experts,
82
+ dtype=torch.float32,
83
+ device=device,
84
+ )
85
+ else:
86
+ self.gemm1_clamp_limit = None
87
+
88
@staticmethod
89
def activation_format() -> mk.FusedMoEActivationFormat:
90
return mk.FusedMoEActivationFormat.Standard
@@ -77,8 +106,12 @@ class TrtLlmFp8ExpertsBase:
106
107
@staticmethod
108
def _supports_activation(activation: MoEActivation) -> bool:
80
- """Supports only SiLU and RELU^2 non-gated activation."""
81
- return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
109
+ """Supports SiLU, SwiGLU-OAI (uninterleaved), and RELU^2 non-gated."""
110
+ return activation in [
111
+ MoEActivation.SILU,
112
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
113
+ MoEActivation.RELU2_NO_MUL,
114
+ ]
115
116
@staticmethod
117
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -198,6 +231,9 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
231
hidden_states_scale=hidden_states_scale,
232
gemm1_weights=w1,
233
gemm1_weights_scale=self.quant_config.w1_scale,
234
+ gemm1_alpha=self.gemm1_alpha,
235
+ gemm1_beta=self.gemm1_beta,
236
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
237
gemm2_weights=w2,
238
gemm2_weights_scale=self.quant_config.w2_scale,
239
num_experts=global_num_experts,
@@ -327,7 +363,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
363
from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
364
365
assert not apply_router_weight_on_input
330
- assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
366
+ assert activation in [
367
+ MoEActivation.SILU,
368
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
369
+ MoEActivation.RELU2_NO_MUL,
370
+ ]
371
activation_type = activation_to_flashinfer_int(activation)
372
assert self.topk <= global_num_experts
373
assert global_num_experts % 4 == 0
@@ -362,6 +402,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
402
hidden_states_scale=hidden_states_scale,
403
gemm1_weights=w1,
404
gemm1_weights_scale=self.quant_config.w1_scale,
405
+ gemm1_alpha=self.gemm1_alpha,
406
+ gemm1_beta=self.gemm1_beta,
407
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
408
gemm2_weights=w2,
409
gemm2_weights_scale=self.quant_config.w2_scale,
410
num_experts=global_num_experts,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
+2
@@ -153,6 +153,8 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
153
a2_scale=layer.w2_input_scale,
154
block_shape=self.weight_block_size,
155
swiglu_limit=getattr(layer, "swiglu_limit", None),
156
+ gemm1_alpha=getattr(layer, "swiglu_alpha", None),
157
+ gemm1_beta=getattr(layer, "swiglu_beta", None),
158
)
159
160
def maybe_make_prepare_finalize(
vllm/model_executor/layers/quantization/online/mxfp8.py
+2
@@ -224,6 +224,8 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
224
w2_bias=getattr(layer, "w2_bias", None),
225
block_shape=self.weight_block_size,
226
swiglu_limit=getattr(layer, "swiglu_limit", None),
227
+ gemm1_alpha=getattr(layer, "swiglu_alpha", None),
228
+ gemm1_beta=getattr(layer, "swiglu_beta", None),
229
)
230
231
def process_weights_after_loading(self, layer: Module) -> None:
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+2
@@ -33,6 +33,8 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
33
MoEActivation.SILU_NO_MUL: ActivationType.Silu,
34
MoEActivation.GELU_NO_MUL: ActivationType.Gelu,
35
MoEActivation.SILU: ActivationType.Swiglu,
36
+ # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args.
37
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
38
MoEActivation.GELU: ActivationType.Geglu,
39
MoEActivation.GELU_TANH: ActivationType.Geglu,
40
MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,