From ef762c26e59302dda6dbbc3620aca86ad2f2bd57 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 15 Jun 2026 18:39:00 +0000 Subject: [PATCH] [MoE] Plumb gemm1_alpha/beta/clamp_limit into TRT-LLM FP8 MoE The FlashInfer FP8 block-scale MoE kernels (trtllm_fp8_block_scale_moe and trtllm_fp8_block_scale_routed_moe) accept optional per-expert SwiGLU parameters gemm1_alpha, gemm1_beta, gemm1_clamp_limit that realize the OAI SwiGLU variant for MXFP8. The FP8 experts did not pass them. - Build per-expert alpha/beta/clamp tensors in TrtLlmFp8ExpertsBase and pass them to the block-scale (monolithic) and routed (modular) kernels, mirroring the existing MXFP4 experts. - Forward gemm1_alpha/gemm1_beta from layer.swiglu_alpha/beta in the MXFP8-capable FP8 quant-config builders (fp8, online mxfp8, compressed-tensors mxfp8). Previously only swiglu_limit was forwarded; only ModelOptMxfp8 forwarded all three, so other paths silently dropped alpha/beta. - Add swigluoai_uninterleave to the supported activations (maps to FlashInfer Swiglu; the OAI behavior comes from the gemm1_* params). Signed-off-by: Yongye Zhu Co-Authored-By: Claude Opus 4.8 (1M context) (cherry picked from commit f4b8540570b7a6a5e0f2e5fd460f4f0943b7cc54) --- .../fused_moe/experts/trtllm_fp8_moe.py | 49 +++++++++++++++++-- .../compressed_tensors_moe_w8a8_mxfp8.py | 2 + .../layers/quantization/online/mxfp8.py | 2 + .../quantization/utils/flashinfer_utils.py | 2 + 4 files changed, 52 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 257bfeee5d3..a7faa5f6e17 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -56,6 +56,35 @@ class TrtLlmFp8ExpertsBase: self.moe_config = moe_config self.quant_config = quant_config + # Per-expert SwiGLU parameters from quant_config (MXFP8 + Swiglu only). + device = torch.accelerator.current_device_index() + if quant_config.gemm1_alpha is not None: + self.gemm1_alpha = torch.tensor( + [quant_config.gemm1_alpha] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_alpha = None + + if quant_config.gemm1_beta is not None: + self.gemm1_beta = torch.tensor( + [quant_config.gemm1_beta] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_beta = None + + if quant_config.gemm1_clamp_limit is not None: + self.gemm1_clamp_limit = torch.tensor( + [quant_config.gemm1_clamp_limit] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_clamp_limit = None + @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard @@ -77,8 +106,12 @@ class TrtLlmFp8ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports only SiLU and RELU^2 non-gated activation.""" - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + """Supports SiLU, SwiGLU-OAI (uninterleaved), and RELU^2 non-gated.""" + return activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.RELU2_NO_MUL, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -198,6 +231,9 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): hidden_states_scale=hidden_states_scale, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, @@ -327,7 +363,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout assert not apply_router_weight_on_input - assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + assert activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.RELU2_NO_MUL, + ] activation_type = activation_to_flashinfer_int(activation) assert self.topk <= global_num_experts assert global_num_experts % 4 == 0 @@ -362,6 +402,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit hidden_states_scale=hidden_states_scale, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py index dc851cc1313..2e6e01ca766 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -153,6 +153,8 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): a2_scale=layer.w2_input_scale, block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index c197398a09b..09d581a0734 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -224,6 +224,8 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): w2_bias=getattr(layer, "w2_bias", None), block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def process_weights_after_loading(self, layer: Module) -> None: diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 1cbfdf69c99..b60deb078e3 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -33,6 +33,8 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" MoEActivation.SILU_NO_MUL: ActivationType.Silu, MoEActivation.GELU_NO_MUL: ActivationType.Gelu, MoEActivation.SILU: ActivationType.Swiglu, + # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args. + MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, MoEActivation.GELU: ActivationType.Geglu, MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,