@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 250 lines 7.95 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5
6 from tests.kernels.moe.utils import make_test_quant_config
7 from tests.kernels.quantization.nvfp4_utils import (
8 FLOAT4_E2M1_MAX,
9 FLOAT8_E4M3_MAX,
10 dequantize_nvfp4_to_dtype,
11 )
12 from tests.kernels.utils import torch_moe
13 from vllm import _custom_ops as ops
14 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
15 from vllm.model_executor.layers.fused_moe import fused_topk
16 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17 from vllm.model_executor.layers.fused_moe.all2all_utils import (
18 maybe_make_prepare_finalize,
19 )
20 from vllm.model_executor.layers.fused_moe.config import (
21 FusedMoEConfig,
22 FusedMoEParallelConfig,
23 FusedMoEQuantConfig,
24 RoutingMethodType,
25 )
26 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
27 FlashInferExperts,
28 is_valid_flashinfer_cutlass_fused_moe,
29 )
30 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
31 from vllm.platforms import current_platform
32 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
33 from vllm.utils.math_utils import next_power_of_2
34 from vllm.utils.torch_utils import set_random_seed
35
36 if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
37 100
38 ):
39 pytest.skip(
40 "Requires flashinfer_cutlass_fused_moe and nvfp4 support",
41 allow_module_level=True,
42 )
43
44 MNK_FACTORS = [
45 (2, 1024, 1024),
46 (2, 3072, 1024),
47 (2, 3072, 1536),
48 (64, 1024, 1536),
49 (64, 3072, 1024),
50 (64, 2048, 1536),
51 (224, 1024, 1024),
52 (224, 1024, 1536),
53 ]
54
55
56 @pytest.mark.parametrize(
57 "activation",
58 [MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE],
59 )
60 def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch):
61 from flashinfer.fused_moe.core import ActivationType
62
63 moe_config = FusedMoEConfig(
64 num_experts=2,
65 experts_per_token=1,
66 hidden_dim=128,
67 intermediate_size=128,
68 num_local_experts=2,
69 num_logical_experts=2,
70 activation=activation,
71 device="cuda",
72 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
73 in_dtype=torch.bfloat16,
74 routing_method=RoutingMethodType.TopK,
75 )
76 quant_config = FusedMoEQuantConfig.make(
77 gemm1_alpha=1.702,
78 gemm1_beta=1.0,
79 gemm1_clamp_limit=7.0,
80 )
81 experts = FlashInferExperts(moe_config=moe_config, quant_config=quant_config)
82
83 call_args = {}
84
85 def fake_flashinfer_cutlass_fused_moe(**kwargs):
86 call_args.update(kwargs)
87
88 monkeypatch.setattr(
89 "vllm.model_executor.layers.fused_moe.experts."
90 "flashinfer_cutlass_moe.flashinfer_cutlass_fused_moe",
91 fake_flashinfer_cutlass_fused_moe,
92 )
93 experts.apply(
94 output=torch.empty((1, 128), device="cuda", dtype=torch.bfloat16),
95 hidden_states=torch.empty((1, 128), device="cuda", dtype=torch.bfloat16),
96 w1=torch.empty((2, 256, 128), device="cuda", dtype=torch.bfloat16),
97 w2=torch.empty((2, 128, 128), device="cuda", dtype=torch.bfloat16),
98 topk_weights=torch.ones((1, 1), device="cuda", dtype=torch.float32),
99 topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64),
100 activation=activation,
101 global_num_experts=2,
102 expert_map=None,
103 a1q_scale=None,
104 a2_scale=None,
105 workspace13=None,
106 workspace2=None,
107 expert_tokens_meta=None,
108 apply_router_weight_on_input=False,
109 )
110
111 assert experts._supports_activation(activation)
112 assert call_args["activation_type"] == ActivationType.Swiglu
113 for name, value in (
114 ("swiglu_alpha", 1.702),
115 ("swiglu_beta", 1.0),
116 ("swiglu_limit", 7.0),
117 ):
118 torch.testing.assert_close(
119 call_args[name],
120 torch.full((2,), value, device="cuda", dtype=torch.float32),
121 )
122
123
124 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
125 @pytest.mark.parametrize("e", [40, 64, 256])
126 @pytest.mark.parametrize("topk", [1, 6, 8])
127 @pytest.mark.parametrize("dtype", [torch.bfloat16])
128 @pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])
129 @torch.inference_mode()
130 def test_flashinfer_fp4_moe_no_graph(
131 m: int,
132 n: int,
133 k: int,
134 e: int,
135 topk: int,
136 dtype: torch.dtype,
137 activation: MoEActivation,
138 workspace_init,
139 ):
140 set_random_seed(7)
141 with set_current_vllm_config(
142 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
143 ):
144 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
145
146 quant_blocksize = 16
147 is_gated_act = activation.is_gated
148
149 w1_q, w2_q, quant_config = make_test_quant_config(
150 e,
151 n,
152 k,
153 in_dtype=dtype,
154 quant_dtype="nvfp4",
155 block_shape=None,
156 per_act_token_quant=False,
157 make_gate=is_gated_act,
158 )
159
160 score = torch.randn((m, e), device="cuda", dtype=dtype)
161 topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
162
163 assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)
164
165 moe_config = FusedMoEConfig(
166 num_experts=e,
167 experts_per_token=topk,
168 hidden_dim=k,
169 intermediate_size=n,
170 num_local_experts=e,
171 num_logical_experts=e,
172 activation=activation,
173 device="cuda",
174 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
175 in_dtype=dtype,
176 routing_method=RoutingMethodType.TopK,
177 max_num_tokens=next_power_of_2(m),
178 )
179
180 flashinfer_experts = FusedMoEKernel(
181 maybe_make_prepare_finalize(
182 moe=moe_config,
183 quant_config=quant_config,
184 allow_new_interface=True,
185 use_monolithic=False,
186 ),
187 FlashInferExperts(moe_config=moe_config, quant_config=quant_config),
188 )
189
190 flashinfer_output = flashinfer_experts.apply(
191 hidden_states=a,
192 w1=w1_q,
193 w2=w2_q,
194 topk_weights=topk_weights,
195 topk_ids=topk_ids,
196 activation=activation,
197 global_num_experts=e,
198 expert_map=None,
199 apply_router_weight_on_input=False,
200 )
201
202 # Reference check:
203 a_global_scale = (
204 (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
205 ).to(torch.float32)
206 a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
207 _, m_k = a_fp4.shape
208 a_in_dtype = dequantize_nvfp4_to_dtype(
209 a_fp4,
210 a_scale_interleaved,
211 a_global_scale,
212 dtype=a.dtype,
213 device=a.device,
214 block_size=quant_blocksize,
215 )
216
217 w1_d = torch.empty(
218 (e, (2 if is_gated_act else 1) * n, k), device="cuda", dtype=dtype
219 )
220 w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
221
222 for idx in range(0, e):
223 w1_d[idx] = dequantize_nvfp4_to_dtype(
224 w1_q[idx],
225 quant_config.w1_scale[idx],
226 (1 / quant_config.g1_alphas[idx]),
227 dtype=dtype,
228 device=w1_q.device,
229 block_size=quant_blocksize,
230 )
231 w2_d[idx] = dequantize_nvfp4_to_dtype(
232 w2_q[idx],
233 quant_config.w2_scale[idx],
234 (1 / quant_config.g2_alphas[idx]),
235 dtype=dtype,
236 device=w2_q.device,
237 block_size=quant_blocksize,
238 )
239
240 torch_output = torch_moe(
241 a_in_dtype, w1_d, w2_d, score, topk, activation=activation
242 )
243
244 torch.testing.assert_close(
245 torch_output, flashinfer_output, atol=1e-1, rtol=1e-1
246 )
247
248
249 if __name__ == "__main__":
250 test_flashinfer_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)