[Kernel] Support Nvfp4 Cutedsl Moe Swiglu-oai and Relu2(non-gated) Activation (#47106)
Signed-off-by: Tiekai Bi <tiekaib@nvidia.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Bi Tiekai committed
Aug 7, 2026 at 08:42 UTC
e08111211bc94f5e7f5ebd0f071ef5300f6f5564
9 files changed
+216
-14
tests/kernels/moe/test_flashinfer_cutedsl_layout.py
new
+43
@@ -0,0 +1,43 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Weight-layout normalization for the FlashInfer CuTeDSL NVFP4 MoE backend."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
10
+ reorder_w13_to_w31_for_flashinfer_cutedsl,
11
+)
12
+
13
+_GATE = torch.tensor([[[1], [2], [3], [4]]])
14
+_UP = torch.tensor([[[10], [20], [30], [40]]])
15
+_EXPECTED = torch.cat([_UP, _GATE], dim=1)
16
+
17
+
18
+def test_reorder_w13_swigluoai_interleaved():
19
+ """gpt-oss w13 is [gate0, up0, gate1, ...] rather than packed [gate; up]."""
20
+ w13 = torch.empty(1, 8, 1, dtype=_GATE.dtype)
21
+ w13[:, 0::2] = _GATE
22
+ w13[:, 1::2] = _UP
23
+
24
+ out, out_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
25
+ MoEActivation.SWIGLUOAI, w13, w13 + 100
26
+ )
27
+
28
+ torch.testing.assert_close(out, _EXPECTED)
29
+ torch.testing.assert_close(out_scale, _EXPECTED + 100)
30
+
31
+
32
+@pytest.mark.parametrize(
33
+ "activation", [MoEActivation.SILU, MoEActivation.SWIGLUOAI_UNINTERLEAVE]
34
+)
35
+def test_reorder_w13_packed_layouts(activation: MoEActivation):
36
+ w13 = torch.cat([_GATE, _UP], dim=1)
37
+
38
+ out, out_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
39
+ activation, w13, w13 + 100
40
+ )
41
+
42
+ torch.testing.assert_close(out, _EXPECTED)
43
+ torch.testing.assert_close(out_scale, _EXPECTED + 100)
tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
+101
-8
@@ -13,9 +13,13 @@ from tests.kernels.quantization.nvfp4_utils import (
13
FLOAT8_E4M3_MAX,
14
break_fp4_bytes,
15
)
16
-from tests.kernels.utils import torch_moe
16
from vllm import _custom_ops as ops
17
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
18
+from vllm.model_executor.layers.activation import (
19
+ SiluAndMul,
20
+ SiluAndMulWithClamp,
21
+ SwigluOAIAndMul,
22
+)
23
from vllm.model_executor.layers.fused_moe import fused_topk
24
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
25
from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -46,6 +50,77 @@ if not has_flashinfer_cutedsl_moe_nvfp4() or not (
50
allow_module_level=True,
51
)
52
53
+_SWIGLU_ALPHA = 1.702
54
+_SWIGLU_BETA = 1.0
55
+# The GEMM1 output is O(0.5) at these tensor scales, so the OAI default of 7.0
56
+# would never clamp anything. Pick a limit that actually bites.
57
+_SWIGLU_LIMIT = 0.3
58
+
59
+_ACT_CASES = [
60
+ pytest.param(MoEActivation.SILU, None, None, None, id="silu"),
61
+ pytest.param(MoEActivation.SILU, None, None, _SWIGLU_LIMIT, id="silu-clamp"),
62
+ pytest.param(MoEActivation.RELU2_NO_MUL, None, None, None, id="relu2_no_mul"),
63
+ pytest.param(
64
+ MoEActivation.SWIGLUOAI,
65
+ _SWIGLU_ALPHA,
66
+ _SWIGLU_BETA,
67
+ _SWIGLU_LIMIT,
68
+ id="swigluoai",
69
+ ),
70
+ pytest.param(
71
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
72
+ _SWIGLU_ALPHA,
73
+ _SWIGLU_BETA,
74
+ _SWIGLU_LIMIT,
75
+ id="swigluoai_uninterleave",
76
+ ),
77
+]
78
+
79
+
80
+def _reference_activation(
81
+ activation: MoEActivation,
82
+ alpha: float | None,
83
+ beta: float | None,
84
+ limit: float | None,
85
+):
86
+ """vLLM's own op for this activation, so the reference is not re-derived."""
87
+ if activation == MoEActivation.RELU2_NO_MUL:
88
+ return lambda x: torch.square(torch.relu(x))
89
+ if activation == MoEActivation.SWIGLUOAI:
90
+ # SwigluOAIAndMul hardcodes beta=1 and reads gate/up interleaved.
91
+ assert beta == 1.0
92
+ return SwigluOAIAndMul(alpha=alpha, limit=limit)
93
+ if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
94
+ return SiluAndMulWithClamp(limit, alpha, beta, compile_native=False)
95
+ if limit is not None:
96
+ return SiluAndMulWithClamp(limit, compile_native=False)
97
+ return SiluAndMul()
98
+
99
+
100
+def _torch_moe_reference(
101
+ a: torch.Tensor,
102
+ w1: torch.Tensor,
103
+ w2: torch.Tensor,
104
+ score: torch.Tensor,
105
+ topk: int,
106
+ act_fn,
107
+) -> torch.Tensor:
108
+ m = a.shape[0]
109
+ weights, ids = torch.topk(torch.softmax(score, dim=-1, dtype=torch.float32), topk)
110
+ x = a.view(m, 1, -1).repeat(1, topk, 1).reshape(m * topk, -1)
111
+ flat_ids = ids.reshape(-1)
112
+ out = torch.zeros(m * topk, w2.shape[1], dtype=a.dtype, device=a.device)
113
+ for expert in range(w1.shape[0]):
114
+ mask = flat_ids == expert
115
+ if mask.any():
116
+ acc = act_fn(x[mask] @ w1[expert].transpose(0, 1)).to(a.dtype)
117
+ out[mask] = acc @ w2[expert].transpose(0, 1)
118
+ return (
119
+ (out.view(m, topk, -1).float() * weights.view(m, topk, 1))
120
+ .sum(dim=1)
121
+ .to(a.dtype)
122
+ )
123
+
124
125
def _quantize_nvfp4_linear(
126
weight: torch.Tensor,
@@ -85,14 +160,19 @@ def _dequantize_nvfp4_linear(
160
161
162
@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)])
163
+@pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES)
164
@pytest.mark.parametrize("dtype", [torch.bfloat16])
165
@torch.inference_mode()
90
-def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
166
+def test_flashinfer_cutedsl_fp4_moe(
167
m: int,
168
n: int,
169
k: int,
170
e: int,
171
topk: int,
172
+ activation: MoEActivation,
173
+ alpha: float | None,
174
+ beta: float | None,
175
+ limit: float | None,
176
dtype: torch.dtype,
177
workspace_init,
178
):
@@ -102,7 +182,8 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
182
):
183
hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
184
105
- w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15
185
+ w1_rows = 2 * n if activation.is_gated else n
186
+ w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15
187
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
188
w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1)
189
w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2)
@@ -112,7 +193,6 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
193
hidden_states, score, topk, renormalize=False
194
)
195
115
- activation = MoEActivation.RELU2_NO_MUL
196
fake_layer = SimpleNamespace(activation=activation)
197
a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
198
a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
@@ -144,6 +224,11 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
224
w1_scale=w1_scale_cutedsl,
225
w2_scale=w2_scale_cutedsl,
226
is_scale_swizzled=False,
227
+ # Unset params must be omitted rather than forwarded as None into
228
+ # the kernel's float-typed SwiGLU arguments.
229
+ gemm1_alpha=alpha,
230
+ gemm1_beta=beta,
231
+ gemm1_clamp_limit=limit,
232
)
233
moe_config = FusedMoEConfig(
234
num_experts=e,
@@ -198,7 +283,7 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
283
dtype=dtype,
284
)
285
201
- w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype)
286
+ w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype)
287
w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
288
for idx in range(e):
289
w1_d[idx] = _dequantize_nvfp4_linear(
@@ -214,17 +299,25 @@ def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
299
dtype=dtype,
300
)
301
217
- torch_output = torch_moe(
302
+ torch_output = _torch_moe_reference(
303
a_in_dtype,
304
w1_d,
305
w2_d,
306
score,
307
topk,
223
- activation=activation,
308
+ _reference_activation(activation, alpha, beta, limit),
309
)
310
torch.testing.assert_close(
311
torch_output,
312
cutedsl_output,
228
- atol=2e-1,
313
+ atol=3e-2,
314
rtol=2e-1,
315
)
316
+ # Outputs here are O(1e-2) while NVFP4 noise is O(1e-3), so an absolute
317
+ # tolerance loose enough for the quantization error also accepts a zero
318
+ # tensor. Compare direction too, which dropped SwiGLU params or a wrong
319
+ # w13 layout would break.
320
+ cosine = torch.nn.functional.cosine_similarity(
321
+ cutedsl_output.flatten().float(), torch_output.flatten().float(), dim=0
322
+ )
323
+ assert cosine > 0.99, f"cosine similarity {cosine:.4f} below 0.99"
vllm/model_executor/layers/fused_moe/config.py
+4
@@ -814,6 +814,8 @@ def nvfp4_moe_quant_config(
814
w1_bias: torch.Tensor | None = None,
815
w2_bias: torch.Tensor | None = None,
816
is_scale_swizzled: bool = True,
817
+ gemm1_alpha: float | None = None,
818
+ gemm1_beta: float | None = None,
819
gemm1_clamp_limit: float | None = None,
820
) -> FusedMoEQuantConfig:
821
"""
@@ -833,6 +835,8 @@ def nvfp4_moe_quant_config(
835
per_out_ch_quant=False,
836
block_shape=None,
837
is_scale_swizzled=is_scale_swizzled,
838
+ gemm1_alpha=gemm1_alpha,
839
+ gemm1_beta=gemm1_beta,
840
gemm1_clamp_limit=gemm1_clamp_limit,
841
)
842
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+26
-1
@@ -59,6 +59,9 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
59
self.global_num_experts = moe_config.num_experts
60
self.ep_rank = moe_config.moe_parallel_config.ep_rank
61
self.local_expert_offset = self.ep_rank * self.local_num_experts
62
+ self.gemm1_alpha = quant_config.gemm1_alpha
63
+ self.gemm1_beta = quant_config.gemm1_beta
64
+ self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
65
66
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
67
layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -93,7 +96,12 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
96
97
@staticmethod
98
def _supports_activation(activation: MoEActivation) -> bool:
96
- return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL)
99
+ return activation in (
100
+ MoEActivation.SILU,
101
+ MoEActivation.SWIGLUOAI,
102
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
103
+ MoEActivation.RELU2_NO_MUL,
104
+ )
105
106
@staticmethod
107
def _supports_parallel_config(
@@ -149,6 +157,22 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
157
# The functional API expects x_sf with trailing dim: (M, K//16, 1).
158
x_sf = a1q_scale.unsqueeze(-1)
159
160
+ # The kernel defaults swiglu_{alpha,beta,limit} to the plain-SwiGLU
161
+ # values, so only forward the ones the model actually sets.
162
+ swiglu_params: dict[str, float | None] = {}
163
+ if activation == MoEActivation.SILU:
164
+ swiglu_params = {"swiglu_limit": self.gemm1_clamp_limit}
165
+ elif activation in (
166
+ MoEActivation.SWIGLUOAI,
167
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
168
+ ):
169
+ swiglu_params = {
170
+ "swiglu_alpha": self.gemm1_alpha,
171
+ "swiglu_beta": self.gemm1_beta,
172
+ "swiglu_limit": self.gemm1_clamp_limit,
173
+ }
174
+ swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None}
175
+
176
flashinfer_cute_dsl_fused_moe_nvfp4(
177
x=hidden_states,
178
x_sf=x_sf,
@@ -167,4 +191,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
191
local_expert_offset=self.local_expert_offset,
192
moe_output=output,
193
activation_type=activation_to_flashinfer_int(activation),
194
+ **swiglu_kwargs,
195
)
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+9
-2
@@ -190,6 +190,7 @@ def select_nvfp4_moe_backend(
190
NVFP4_BACKENDS_WITH_CLAMP = {
191
NvFp4MoeBackend.FLASHINFER_TRTLLM,
192
NvFp4MoeBackend.FLASHINFER_CUTLASS,
193
+ NvFp4MoeBackend.FLASHINFER_CUTEDSL,
194
NvFp4MoeBackend.VLLM_CUTLASS,
195
NvFp4MoeBackend.MARLIN,
196
NvFp4MoeBackend.HUMMING,
@@ -261,8 +262,8 @@ def select_nvfp4_moe_backend(
262
f"Model sets swiglu_limit={config.swiglu_limit}, but the "
263
f"explicitly requested moe_backend={runner_backend!r} does "
264
f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', "
264
- f"'flashinfer_cutlass', 'cutlass', 'marlin', or 'humming' "
265
- f"instead."
265
+ f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', "
266
+ f"'marlin', or 'humming' instead."
267
)
268
return _return_or_raise(
269
requested_backend, config, weight_key, activation_key, activation_format
@@ -471,6 +472,8 @@ def make_nvfp4_moe_quant_config(
472
a13_scale: torch.Tensor,
473
a2_scale: torch.Tensor,
474
swiglu_limit: float | None = None,
475
+ swiglu_alpha: float | None = None,
476
+ swiglu_beta: float | None = None,
477
layer: torch.nn.Module | None = None,
478
) -> FusedMoEQuantConfig:
479
if backend == NvFp4MoeBackend.HUMMING:
@@ -497,6 +500,8 @@ def make_nvfp4_moe_quant_config(
500
a2_gscale=a2_scale,
501
w1_scale=w13_scale,
502
w2_scale=w2_scale,
503
+ gemm1_alpha=swiglu_alpha,
504
+ gemm1_beta=swiglu_beta,
505
gemm1_clamp_limit=swiglu_limit,
506
)
507
@@ -521,6 +526,8 @@ def make_nvfp4_moe_quant_config(
526
NvFp4MoeBackend.FLASHINFER_CUTEDSL,
527
)
528
),
529
+ gemm1_alpha=swiglu_alpha,
530
+ gemm1_beta=swiglu_beta,
531
gemm1_clamp_limit=swiglu_limit,
532
)
533
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+2
@@ -251,6 +251,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
251
a13_scale=layer.w13_input_scale,
252
a2_scale=layer.w2_input_scale,
253
swiglu_limit=getattr(layer, "swiglu_limit", None),
254
+ swiglu_alpha=getattr(layer, "swiglu_alpha", None),
255
+ swiglu_beta=getattr(layer, "swiglu_beta", None),
256
layer=layer,
257
)
258
vllm/model_executor/layers/quantization/modelopt.py
+2
@@ -1591,6 +1591,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1591
a13_scale=layer.w13_input_scale,
1592
a2_scale=layer.w2_input_scale,
1593
swiglu_limit=getattr(layer, "swiglu_limit", None),
1594
+ swiglu_alpha=getattr(layer, "swiglu_alpha", None),
1595
+ swiglu_beta=getattr(layer, "swiglu_beta", None),
1596
layer=layer,
1597
)
1598
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3
@@ -1659,6 +1659,9 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1659
w2_scale_2=layer.w2_weight_scale_2,
1660
a13_scale=layer.w13_input_scale_2,
1661
a2_scale=layer.w2_input_scale_2,
1662
+ swiglu_limit=getattr(layer, "swiglu_limit", None),
1663
+ swiglu_alpha=getattr(layer, "swiglu_alpha", None),
1664
+ swiglu_beta=getattr(layer, "swiglu_beta", None),
1665
layer=layer,
1666
)
1667
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+26
-3
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
7
import torch
8
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
11
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
12
align_fp4_moe_weights_for_fi,
13
align_trtllm_fp4_moe_hidden_dim_for_fi,
@@ -80,6 +81,28 @@ def interleave_linear_and_gate(
81
return x
82
83
84
+def reorder_w13_to_w31_for_flashinfer_cutedsl(
85
+ activation: MoEActivation,
86
+ w13: torch.Tensor,
87
+ w13_scale: torch.Tensor,
88
+) -> tuple[torch.Tensor, torch.Tensor]:
89
+ """Normalize gated w13 rows to the [up; gate] order used by FlashInfer."""
90
+ if activation == MoEActivation.SWIGLUOAI:
91
+ # gpt-oss checkpoints store w13 interleaved as [gate0, up0, gate1, ...].
92
+ gate, up = w13[:, 0::2], w13[:, 1::2]
93
+ gate_scale, up_scale = w13_scale[:, 0::2], w13_scale[:, 1::2]
94
+ return (
95
+ torch.cat([up, gate], dim=1).contiguous(),
96
+ torch.cat([up_scale, gate_scale], dim=1).contiguous(),
97
+ )
98
+
99
+ half = w13.shape[1] // 2
100
+ return (
101
+ torch.cat([w13[:, half:], w13[:, :half]], dim=1).contiguous(),
102
+ torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1).contiguous(),
103
+ )
104
+
105
+
106
def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
107
layer: "RoutedExperts",
108
w13: torch.Tensor,
@@ -114,9 +137,9 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
137
a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
138
139
if layer.activation.is_gated:
117
- half = w13.shape[1] // 2
118
- w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1)
119
- w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1)
140
+ w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
141
+ layer.activation, w13, w13_scale
142
+ )
143
144
# Interleave up/gate rows for w13 weights and scales.
145
w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)