Integrate CuTeDSL MoE for ReLU2 NVFP4 (#49580)
Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
danielafrimi committed
Jul 29, 2026 at 05:46 UTC
32a423ac0aad67f94f93e97f73338f484b55faec
3 files changed
+245
-9
tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
new
+230
@@ -0,0 +1,230 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for FlashInfer CuTeDSL NVFP4 MoE."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from tests.kernels.quantization.nvfp4_utils import (
12
+ FLOAT4_E2M1_MAX,
13
+ FLOAT8_E4M3_MAX,
14
+ break_fp4_bytes,
15
+)
16
+from tests.kernels.utils import torch_moe
17
+from vllm import _custom_ops as ops
18
+from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
19
+from vllm.model_executor.layers.fused_moe import fused_topk
20
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
21
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
22
+ maybe_make_prepare_finalize,
23
+)
24
+from vllm.model_executor.layers.fused_moe.config import (
25
+ FusedMoEConfig,
26
+ FusedMoEParallelConfig,
27
+ RoutingMethodType,
28
+ nvfp4_moe_quant_config,
29
+)
30
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
31
+ FlashInferCuteDSLExperts,
32
+)
33
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
34
+ prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
35
+)
36
+from vllm.platforms import current_platform
37
+from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4
38
+from vllm.utils.math_utils import next_power_of_2
39
+from vllm.utils.torch_utils import set_random_seed
40
+
41
+if not has_flashinfer_cutedsl_moe_nvfp4() or not (
42
+ current_platform.is_device_capability_family(100)
43
+):
44
+ pytest.skip(
45
+ "Requires FlashInfer CuTeDSL NVFP4 MoE on SM100",
46
+ allow_module_level=True,
47
+ )
48
+
49
+
50
+def _quantize_nvfp4_linear(
51
+ weight: torch.Tensor,
52
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
53
+ weights_q = []
54
+ scales = []
55
+ global_scales = []
56
+ for expert_weight in weight:
57
+ global_scale = (
58
+ FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max()
59
+ ).to(torch.float32)
60
+ weight_q, scale = ops.scaled_fp4_quant(
61
+ expert_weight,
62
+ global_scale,
63
+ is_sf_swizzled_layout=False,
64
+ )
65
+ weights_q.append(weight_q)
66
+ scales.append(scale)
67
+ global_scales.append(global_scale)
68
+ return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales)
69
+
70
+
71
+def _dequantize_nvfp4_linear(
72
+ tensor_fp4: torch.Tensor,
73
+ tensor_sf: torch.Tensor,
74
+ global_scale: torch.Tensor,
75
+ dtype: torch.dtype,
76
+) -> torch.Tensor:
77
+ assert tensor_fp4.dtype == torch.uint8
78
+ m, packed_k = tensor_fp4.shape
79
+ k = packed_k * 2
80
+ tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
81
+ tensor_f32 = tensor_f32.reshape(m, k // 16, 16)
82
+ tensor_sf = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32)
83
+ tensor_sf = tensor_sf[:, : k // 16] / global_scale
84
+ return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype)
85
+
86
+
87
+@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)])
88
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
89
+@torch.inference_mode()
90
+def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
91
+ m: int,
92
+ n: int,
93
+ k: int,
94
+ e: int,
95
+ topk: int,
96
+ dtype: torch.dtype,
97
+ workspace_init,
98
+):
99
+ set_random_seed(7)
100
+ with set_current_vllm_config(
101
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
102
+ ):
103
+ hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
104
+
105
+ w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15
106
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
107
+ w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1)
108
+ w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2)
109
+
110
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
111
+ topk_weights, topk_ids, _ = fused_topk(
112
+ hidden_states, score, topk, renormalize=False
113
+ )
114
+
115
+ activation = MoEActivation.RELU2_NO_MUL
116
+ fake_layer = SimpleNamespace(activation=activation)
117
+ a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
118
+ a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
119
+ (
120
+ w1_cutedsl,
121
+ w1_scale_cutedsl,
122
+ w1_alpha,
123
+ a1_scale,
124
+ w2_cutedsl,
125
+ w2_scale_cutedsl,
126
+ w2_alpha,
127
+ a2_scale,
128
+ ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
129
+ layer=fake_layer,
130
+ w13=w1_q,
131
+ w13_scale=w1_scale,
132
+ w13_scale_2=(1.0 / w1_global_scale),
133
+ a13_scale=a1_scale,
134
+ w2=w2_q,
135
+ w2_scale=w2_scale,
136
+ w2_scale_2=(1.0 / w2_global_scale),
137
+ a2_scale=a2_scale,
138
+ )
139
+ quant_config = nvfp4_moe_quant_config(
140
+ g1_alphas=w1_alpha,
141
+ g2_alphas=w2_alpha,
142
+ a1_gscale=(1.0 / a1_scale),
143
+ a2_gscale=(1.0 / a2_scale),
144
+ w1_scale=w1_scale_cutedsl,
145
+ w2_scale=w2_scale_cutedsl,
146
+ is_scale_swizzled=False,
147
+ )
148
+ moe_config = FusedMoEConfig(
149
+ num_experts=e,
150
+ experts_per_token=topk,
151
+ hidden_dim=k,
152
+ intermediate_size=n,
153
+ num_local_experts=e,
154
+ num_logical_experts=e,
155
+ activation=activation,
156
+ device="cuda",
157
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
158
+ in_dtype=dtype,
159
+ routing_method=RoutingMethodType.TopK,
160
+ max_num_tokens=next_power_of_2(m),
161
+ )
162
+
163
+ cutedsl_experts = mk.FusedMoEKernel(
164
+ maybe_make_prepare_finalize(
165
+ moe=moe_config,
166
+ quant_config=quant_config,
167
+ allow_new_interface=True,
168
+ use_monolithic=False,
169
+ ),
170
+ FlashInferCuteDSLExperts(
171
+ moe_config=moe_config,
172
+ quant_config=quant_config,
173
+ ),
174
+ )
175
+
176
+ cutedsl_output = cutedsl_experts.apply(
177
+ hidden_states=hidden_states,
178
+ w1=w1_cutedsl,
179
+ w2=w2_cutedsl,
180
+ topk_weights=topk_weights,
181
+ topk_ids=topk_ids,
182
+ activation=activation,
183
+ global_num_experts=e,
184
+ expert_map=None,
185
+ apply_router_weight_on_input=False,
186
+ )
187
+
188
+ a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32)
189
+ a_q, a_scale = ops.scaled_fp4_quant(
190
+ hidden_states,
191
+ a_global_scale,
192
+ is_sf_swizzled_layout=False,
193
+ )
194
+ a_in_dtype = _dequantize_nvfp4_linear(
195
+ a_q,
196
+ a_scale,
197
+ a_global_scale,
198
+ dtype=dtype,
199
+ )
200
+
201
+ w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype)
202
+ w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
203
+ for idx in range(e):
204
+ w1_d[idx] = _dequantize_nvfp4_linear(
205
+ w1_q[idx],
206
+ w1_scale[idx],
207
+ w1_global_scale[idx],
208
+ dtype=dtype,
209
+ )
210
+ w2_d[idx] = _dequantize_nvfp4_linear(
211
+ w2_q[idx],
212
+ w2_scale[idx],
213
+ w2_global_scale[idx],
214
+ dtype=dtype,
215
+ )
216
+
217
+ torch_output = torch_moe(
218
+ a_in_dtype,
219
+ w1_d,
220
+ w2_d,
221
+ score,
222
+ topk,
223
+ activation=activation,
224
+ )
225
+ torch.testing.assert_close(
226
+ torch_output,
227
+ cutedsl_output,
228
+ atol=2e-1,
229
+ rtol=2e-1,
230
+ )
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+6
-2
@@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import (
13
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
TopKWeightAndReduceNoOP,
15
)
16
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
17
+ activation_to_flashinfer_int,
18
+)
19
from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
QuantKey,
21
kNvfp4Dynamic,
@@ -76,7 +79,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
79
80
@staticmethod
81
def _supports_no_act_and_mul() -> bool:
79
- return False
82
+ return True
83
84
@staticmethod
85
def _supports_quant_scheme(
@@ -90,7 +93,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
93
94
@staticmethod
95
def _supports_activation(activation: MoEActivation) -> bool:
93
- return activation == MoEActivation.SILU
96
+ return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL)
97
98
@staticmethod
99
def _supports_parallel_config(
@@ -163,4 +166,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
166
num_local_experts=self.local_num_experts,
167
local_expert_offset=self.local_expert_offset,
168
moe_output=output,
169
+ activation_type=activation_to_flashinfer_int(activation),
170
)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+9
-7
@@ -103,7 +103,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
103
"""Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.
104
105
Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper,
106
- and interleaves w13 gate/linear rows.
106
+ and interleaves w13 gate/linear rows for gated activations. Non-gated
107
+ activations use a single w13 projection and keep its row order unchanged.
108
"""
109
from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
110
@@ -112,13 +113,14 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
113
a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
114
a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
115
115
- half = w13.shape[1] // 2
116
- w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1)
117
- w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1)
116
+ 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)
120
119
- # Interleave up/gate rows for w13 weights and scales.
120
- w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
121
- w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
121
+ # Interleave up/gate rows for w13 weights and scales.
122
+ w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
123
+ w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
124
125
# Convert w13 scale factors: linear → swizzled → MMA layout.
126
w13_scale = swizzle_blockscale(w13_scale)