[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>
aoshen02 committed
Sep 4, 2026 at 23:34 UTC
f19431e5c8634cd0d1641d3f7f4a01043f0eddab
2 files changed
+84
-6
tests/kernels/moe/test_flashinfer.py
+70
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from tests.kernels.moe.utils import make_dummy_moe_config
10
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12
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
23
FlashInferExperts,
24
)
25
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
26
+ TrtLlmFp8ExpertsModular,
27
TrtLlmFp8ExpertsMonolithic,
28
)
29
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 (
32
swap_w13_to_w31,
33
)
34
from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
35
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
+ QuantKey,
37
+ kFp8Dynamic128Sym,
38
+ kFp8DynamicTensorSym,
39
+ kFp8Static128BlockSym,
40
+ kFp8StaticTensorSym,
41
+ kMxfp8Dynamic,
42
+ kMxfp8Static,
43
+)
44
from vllm.model_executor.models.llama4 import Llama4MoE
45
from vllm.platforms import current_platform
46
from vllm.utils.math_utils import next_power_of_2
@@ -422,3 +433,62 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
433
434
assert w13_converted.shape[0] == num_experts
435
assert w2_converted.shape[0] == num_experts
436
+
437
+
438
+@pytest.mark.parametrize(
439
+ ("weight_key", "activation_key", "activation", "expected"),
440
+ [
441
+ (kMxfp8Static, kMxfp8Dynamic, MoEActivation.SILU, True),
442
+ (kFp8Static128BlockSym, kFp8Dynamic128Sym, MoEActivation.SILU, True),
443
+ (
444
+ kFp8Static128BlockSym,
445
+ kFp8Dynamic128Sym,
446
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
447
+ True,
448
+ ),
449
+ (kFp8StaticTensorSym, kFp8DynamicTensorSym, MoEActivation.SILU, False),
450
+ # FlashInfer takes the clamp only with a SwiGLU activation.
451
+ (
452
+ kFp8Static128BlockSym,
453
+ kFp8Dynamic128Sym,
454
+ MoEActivation.RELU2_NO_MUL,
455
+ False,
456
+ ),
457
+ ],
458
+)
459
+def test_trtllm_fp8_swiglu_clamp_support(
460
+ weight_key: QuantKey,
461
+ activation_key: QuantKey,
462
+ activation: MoEActivation,
463
+ expected: bool,
464
+):
465
+ """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled
466
+ kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the
467
+ per-tensor kernel has no clamp, and Relu2 rejects the parameters."""
468
+
469
+ class _Experts(TrtLlmFp8ExpertsModular):
470
+ @staticmethod
471
+ def _supports_current_device() -> bool:
472
+ return True
473
+
474
+ @staticmethod
475
+ def _supports_quant_scheme(
476
+ weight_key: QuantKey | None, activation_key: QuantKey | None
477
+ ) -> bool:
478
+ return True
479
+
480
+ moe_config = make_dummy_moe_config()
481
+ moe_config.swiglu_limit = 7.0
482
+ moe_config.activation = activation
483
+
484
+ supported, reason = _Experts.is_supported_config(
485
+ _Experts,
486
+ moe_config,
487
+ weight_key,
488
+ activation_key,
489
+ mk.FusedMoEActivationFormat.Standard,
490
+ )
491
+
492
+ assert supported == expected, reason
493
+ if not expected:
494
+ assert "SwiGLU" in reason
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+14
-6
@@ -80,18 +80,26 @@ class TrtLlmFp8ExpertsBase:
80
)
81
if not supported:
82
return supported, reason
83
- # FlashInfer accepts gemm1_alpha/beta/clamp_limit only for MXFP8; the
84
- # DeepSeek-FP8 block-scale kernel rejects them and the per-tensor kernel
85
- # has no clamp at all, so a model-requested SwiGLU clamp cannot be
86
- # honored on those paths.
83
if (
84
moe_config.swiglu_limit is not None
85
or moe_config.swiglu_alpha is not None
86
or moe_config.swiglu_beta is not None
91
- ) and (weight_key, activation_key) != (kMxfp8Static, kMxfp8Dynamic):
87
+ ) and (
88
+ (weight_key, activation_key)
89
+ not in (
90
+ (kMxfp8Static, kMxfp8Dynamic),
91
+ (kFp8Static128BlockSym, kFp8Dynamic128Sym),
92
+ )
93
+ or moe_config.activation
94
+ not in (
95
+ MoEActivation.SILU,
96
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
97
+ )
98
+ ):
99
return False, (
100
"the TRTLLM FP8 kernels apply the SwiGLU alpha/beta/clamp "
94
- "parameters only for MXFP8 weights"
101
+ "parameters only for block-scaled weights with a SwiGLU "
102
+ f"activation, but got {weight_key} and {moe_config.activation}"
103
)
104
if moe_config.num_experts > 2048:
105
return False, (