@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 357 lines 12.4 KB
Raw
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 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 (
26 maybe_make_prepare_finalize,
27 )
28 from vllm.model_executor.layers.fused_moe.config import (
29 FusedMoEConfig,
30 FusedMoEParallelConfig,
31 RoutingMethodType,
32 )
33 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
34 FlashInferCuteDSLExperts,
35 )
36 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
37 NvFp4MoeBackend,
38 make_nvfp4_moe_quant_config,
39 )
40 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
41 prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
42 )
43 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
44 quantize_nvfp4_per_token_input,
45 )
46 from vllm.platforms import current_platform
47 from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4
48 from vllm.utils.math_utils import next_power_of_2
49 from vllm.utils.torch_utils import set_random_seed
50
51 if not has_flashinfer_cutedsl_moe_nvfp4() or not (
52 current_platform.is_device_capability_family(100)
53 ):
54 pytest.skip(
55 "Requires FlashInfer CuTeDSL NVFP4 MoE on SM100",
56 allow_module_level=True,
57 )
58
59 _SWIGLU_ALPHA = 1.702
60 _SWIGLU_BETA = 1.0
61 # The GEMM1 output is O(0.5) at these tensor scales, so the OAI default of 7.0
62 # would never clamp anything. Pick a limit that actually bites.
63 _SWIGLU_LIMIT = 0.3
64
65 _ACT_CASES = [
66 pytest.param(MoEActivation.SILU, None, None, None, id="silu"),
67 pytest.param(MoEActivation.SILU, None, None, _SWIGLU_LIMIT, id="silu-clamp"),
68 pytest.param(MoEActivation.RELU2_NO_MUL, None, None, None, id="relu2_no_mul"),
69 pytest.param(
70 MoEActivation.SWIGLUOAI,
71 _SWIGLU_ALPHA,
72 _SWIGLU_BETA,
73 _SWIGLU_LIMIT,
74 id="swigluoai",
75 ),
76 pytest.param(
77 MoEActivation.SWIGLUOAI_UNINTERLEAVE,
78 _SWIGLU_ALPHA,
79 _SWIGLU_BETA,
80 _SWIGLU_LIMIT,
81 id="swigluoai_uninterleave",
82 ),
83 ]
84
85
86 def _reference_activation(
87 activation: MoEActivation,
88 alpha: float | None,
89 beta: float | None,
90 limit: float | None,
91 ):
92 """VLLM's own op for this activation, so the reference is not re-derived."""
93 if activation == MoEActivation.RELU2_NO_MUL:
94 return lambda x: torch.square(torch.relu(x))
95 if activation == MoEActivation.SWIGLUOAI:
96 # SwigluOAIAndMul hardcodes beta=1 and reads gate/up interleaved.
97 assert beta == 1.0
98 return SwigluOAIAndMul(alpha=alpha, limit=limit)
99 if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
100 return SiluAndMulWithClamp(limit, alpha, beta, compile_native=False)
101 if limit is not None:
102 return SiluAndMulWithClamp(limit, compile_native=False)
103 return SiluAndMul()
104
105
106 def _torch_moe_reference(
107 a: torch.Tensor,
108 w1: torch.Tensor,
109 w2: torch.Tensor,
110 score: torch.Tensor,
111 topk: int,
112 act_fn,
113 ) -> torch.Tensor:
114 m = a.shape[0]
115 weights, ids = torch.topk(torch.softmax(score, dim=-1, dtype=torch.float32), topk)
116 x = a.view(m, 1, -1).repeat(1, topk, 1).reshape(m * topk, -1)
117 flat_ids = ids.reshape(-1)
118 out = torch.zeros(m * topk, w2.shape[1], dtype=a.dtype, device=a.device)
119 for expert in range(w1.shape[0]):
120 mask = flat_ids == expert
121 if mask.any():
122 acc = act_fn(x[mask] @ w1[expert].transpose(0, 1)).to(a.dtype)
123 out[mask] = acc @ w2[expert].transpose(0, 1)
124 return (
125 (out.view(m, topk, -1).float() * weights.view(m, topk, 1))
126 .sum(dim=1)
127 .to(a.dtype)
128 )
129
130
131 def _quantize_nvfp4_linear(
132 weight: torch.Tensor,
133 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
134 weights_q = []
135 scales = []
136 global_scales = []
137 for expert_weight in weight:
138 global_scale = (
139 FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max()
140 ).to(torch.float32)
141 weight_q, scale = ops.scaled_fp4_quant(
142 expert_weight,
143 global_scale,
144 is_sf_swizzled_layout=False,
145 )
146 weights_q.append(weight_q)
147 scales.append(scale)
148 global_scales.append(global_scale)
149 return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales)
150
151
152 def _dequantize_nvfp4_linear(
153 tensor_fp4: torch.Tensor,
154 tensor_sf: torch.Tensor,
155 global_scale: torch.Tensor,
156 dtype: torch.dtype,
157 ) -> torch.Tensor:
158 assert tensor_fp4.dtype == torch.uint8
159 m, packed_k = tensor_fp4.shape
160 k = packed_k * 2
161 tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
162 tensor_f32 = tensor_f32.reshape(m, k // 16, 16)
163 tensor_sf = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32)
164 tensor_sf = tensor_sf[:, : k // 16] / global_scale
165 return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype)
166
167
168 @pytest.mark.parametrize(
169 "m,n,k,e,topk",
170 [(16, 128, 512, 4, 2), (16, 192, 512, 4, 2), (16, 192, 2688, 4, 2)],
171 )
172 @pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES)
173 @pytest.mark.parametrize("per_token_activation", [False, True])
174 @pytest.mark.parametrize("dtype", [torch.bfloat16])
175 @torch.inference_mode()
176 def test_flashinfer_cutedsl_fp4_moe(
177 m: int,
178 n: int,
179 k: int,
180 e: int,
181 topk: int,
182 activation: MoEActivation,
183 alpha: float | None,
184 beta: float | None,
185 limit: float | None,
186 per_token_activation: bool,
187 dtype: torch.dtype,
188 workspace_init,
189 ):
190 set_random_seed(7)
191 with set_current_vllm_config(
192 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
193 ):
194 hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
195 # Keep GEMM1 activation variance stable as K changes, so the fixed
196 # tolerance against the BF16 GEMM2 reference measures comparable FP4
197 # quantization error at each hidden width.
198 hidden_states *= (512 / k) ** 0.5
199
200 w1_rows = 2 * n if activation.is_gated else n
201 w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15
202 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
203 w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1)
204 w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2)
205
206 score = torch.randn((m, e), device="cuda", dtype=dtype)
207 topk_weights, topk_ids, _ = fused_topk(
208 hidden_states, score, topk, renormalize=False
209 )
210
211 moe_config = FusedMoEConfig(
212 num_experts=e,
213 experts_per_token=topk,
214 hidden_dim=k,
215 intermediate_size=n,
216 num_local_experts=e,
217 num_logical_experts=e,
218 activation=activation,
219 device="cuda",
220 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
221 in_dtype=dtype,
222 routing_method=RoutingMethodType.TopK,
223 max_num_tokens=next_power_of_2(m),
224 )
225 fake_layer = SimpleNamespace(activation=activation, moe_config=moe_config)
226 a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
227 a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
228 (
229 w1_cutedsl,
230 w1_scale_cutedsl,
231 w1_alpha,
232 a1_scale,
233 w2_cutedsl,
234 w2_scale_cutedsl,
235 w2_alpha,
236 a2_scale,
237 ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
238 layer=fake_layer,
239 w13=w1_q,
240 w13_scale=w1_scale,
241 w13_scale_2=(1.0 / w1_global_scale),
242 a13_scale=a1_scale,
243 w2=w2_q,
244 w2_scale=w2_scale,
245 w2_scale_2=(1.0 / w2_global_scale),
246 a2_scale=a2_scale,
247 )
248 # I=192 must remain unpadded for gated experts, but become 256 for
249 # non-gated experts. The reference below uses the checkpoint shapes.
250 expected_n = n if activation.is_gated else (n + 127) // 128 * 128
251 assert moe_config.intermediate_size_per_partition == expected_n
252 assert w1_cutedsl.shape[1] == expected_n * (2 if activation.is_gated else 1)
253 assert w2_cutedsl.shape[2] * 2 == expected_n
254 expected_k = (k + 255) // 256 * 256
255 assert moe_config.hidden_dim == expected_k
256 assert moe_config.hidden_dim_unpadded == k
257 assert w1_cutedsl.shape[2] * 2 == expected_k
258 assert w2_cutedsl.shape[1] == expected_k
259 quant_config = make_nvfp4_moe_quant_config(
260 backend=NvFp4MoeBackend.FLASHINFER_CUTEDSL,
261 w13_scale_2=w1_alpha,
262 w2_scale_2=w2_alpha,
263 a13_scale=a1_scale,
264 a2_scale=a2_scale,
265 w13_scale=w1_scale_cutedsl,
266 w2_scale=w2_scale_cutedsl,
267 # Unset params must be omitted rather than forwarded as None into
268 # the kernel's float-typed SwiGLU arguments.
269 swiglu_alpha=alpha,
270 swiglu_beta=beta,
271 swiglu_limit=limit,
272 )
273
274 cutedsl_experts = mk.FusedMoEKernel(
275 maybe_make_prepare_finalize(
276 moe=moe_config,
277 quant_config=quant_config,
278 allow_new_interface=True,
279 use_monolithic=False,
280 ),
281 FlashInferCuteDSLExperts(
282 moe_config=moe_config,
283 quant_config=quant_config,
284 per_token_activation=per_token_activation,
285 ),
286 )
287
288 cutedsl_output = cutedsl_experts.apply(
289 # The model runner pads inputs and trims outputs around this kernel.
290 hidden_states=torch.nn.functional.pad(hidden_states, (0, expected_k - k)),
291 w1=w1_cutedsl,
292 w2=w2_cutedsl,
293 topk_weights=topk_weights,
294 topk_ids=topk_ids,
295 activation=activation,
296 global_num_experts=e,
297 expert_map=None,
298 apply_router_weight_on_input=False,
299 )
300 cutedsl_output = cutedsl_output[:, :k]
301
302 a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32)
303 if per_token_activation:
304 a_q, a_scale, a_row_scale = quantize_nvfp4_per_token_input(hidden_states)
305 else:
306 a_q, a_scale = ops.scaled_fp4_quant(
307 hidden_states,
308 a_global_scale,
309 is_sf_swizzled_layout=False,
310 )
311 a_in_dtype = _dequantize_nvfp4_linear(
312 a_q,
313 a_scale,
314 a_global_scale,
315 dtype=dtype,
316 )
317 if per_token_activation:
318 a_in_dtype = (a_in_dtype.float() * a_row_scale.view(-1, 1)).to(dtype)
319
320 w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype)
321 w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
322 for idx in range(e):
323 w1_d[idx] = _dequantize_nvfp4_linear(
324 w1_q[idx],
325 w1_scale[idx],
326 w1_global_scale[idx],
327 dtype=dtype,
328 )
329 w2_d[idx] = _dequantize_nvfp4_linear(
330 w2_q[idx],
331 w2_scale[idx],
332 w2_global_scale[idx],
333 dtype=dtype,
334 )
335
336 torch_output = _torch_moe_reference(
337 a_in_dtype,
338 w1_d,
339 w2_d,
340 score,
341 topk,
342 _reference_activation(activation, alpha, beta, limit),
343 )
344 torch.testing.assert_close(
345 torch_output,
346 cutedsl_output,
347 atol=3e-2,
348 rtol=2e-1,
349 )
350 # Outputs here are O(1e-2) while NVFP4 noise is O(1e-3), so an absolute
351 # tolerance loose enough for the quantization error also accepts a zero
352 # tensor. Compare direction too, which dropped SwiGLU params or a wrong
353 # w13 layout would break.
354 cosine = torch.nn.functional.cosine_similarity(
355 cutedsl_output.flatten().float(), torch_output.flatten().float(), dim=0
356 )
357 assert cosine > 0.99, f"cosine similarity {cosine:.4f} below 0.99"