[Bugfix][MoE] Allow TRTLLM FP8 block-scale MoE with SwiGLU clamp (#55069)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen524@gmail.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>

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

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, (