diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index b1cfa51190..2e022ede36 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -6,6 +6,7 @@ import pytest import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.moe.utils import make_dummy_moe_config from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -22,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import FlashInferExperts, ) from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8ExpertsModular, TrtLlmFp8ExpertsMonolithic, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts @@ -30,6 +32,15 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8DynamicTensorSym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, + kMxfp8Dynamic, + kMxfp8Static, +) from vllm.model_executor.models.llama4 import Llama4MoE from vllm.platforms import current_platform from vllm.utils.math_utils import next_power_of_2 @@ -422,3 +433,62 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout( assert w13_converted.shape[0] == num_experts assert w2_converted.shape[0] == num_experts + + +@pytest.mark.parametrize( + ("weight_key", "activation_key", "activation", "expected"), + [ + (kMxfp8Static, kMxfp8Dynamic, MoEActivation.SILU, True), + (kFp8Static128BlockSym, kFp8Dynamic128Sym, MoEActivation.SILU, True), + ( + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + True, + ), + (kFp8StaticTensorSym, kFp8DynamicTensorSym, MoEActivation.SILU, False), + # FlashInfer takes the clamp only with a SwiGLU activation. + ( + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + MoEActivation.RELU2_NO_MUL, + False, + ), + ], +) +def test_trtllm_fp8_swiglu_clamp_support( + weight_key: QuantKey, + activation_key: QuantKey, + activation: MoEActivation, + expected: bool, +): + """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled + kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the + per-tensor kernel has no clamp, and Relu2 rejects the parameters.""" + + class _Experts(TrtLlmFp8ExpertsModular): + @staticmethod + def _supports_current_device() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, activation_key: QuantKey | None + ) -> bool: + return True + + moe_config = make_dummy_moe_config() + moe_config.swiglu_limit = 7.0 + moe_config.activation = activation + + supported, reason = _Experts.is_supported_config( + _Experts, + moe_config, + weight_key, + activation_key, + mk.FusedMoEActivationFormat.Standard, + ) + + assert supported == expected, reason + if not expected: + assert "SwiGLU" in reason 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 709ba17d42..6f256db960 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 @@ -80,18 +80,26 @@ class TrtLlmFp8ExpertsBase: ) if not supported: return supported, reason - # FlashInfer accepts gemm1_alpha/beta/clamp_limit only for MXFP8; the - # DeepSeek-FP8 block-scale kernel rejects them and the per-tensor kernel - # has no clamp at all, so a model-requested SwiGLU clamp cannot be - # honored on those paths. if ( moe_config.swiglu_limit is not None or moe_config.swiglu_alpha is not None or moe_config.swiglu_beta is not None - ) and (weight_key, activation_key) != (kMxfp8Static, kMxfp8Dynamic): + ) and ( + (weight_key, activation_key) + not in ( + (kMxfp8Static, kMxfp8Dynamic), + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + ) + or moe_config.activation + not in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + ): return False, ( "the TRTLLM FP8 kernels apply the SwiGLU alpha/beta/clamp " - "parameters only for MXFP8 weights" + "parameters only for block-scaled weights with a SwiGLU " + f"activation, but got {weight_key} and {moe_config.activation}" ) if moe_config.num_experts > 2048: return False, (