@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 211 lines 6.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for ZenTorch CPU fused MoE dispatch and forward."""
4
5 import pytest
6 import torch
7
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
10 from tests.kernels.moe.test_cpu_fused_moe import (
11 UNALIGNED_INTERMEDIATE_DIM,
12 _make_moe_config,
13 _StubMoELayer,
14 ref_fused_moe,
15 )
16 from vllm.model_executor.kernels.linear.zentorch_utils import (
17 has_zentorch_op,
18 is_zentorch_moe_config_supported,
19 )
20 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
21 from vllm.model_executor.layers.fused_moe.config import (
22 FusedMoEConfig,
23 FusedMoEQuantConfig,
24 biased_moe_quant_config,
25 )
26 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
27 CPUUnquantizedExperts,
28 )
29 from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk
30 from vllm.platforms import current_platform
31 from vllm.utils.torch_utils import set_random_seed
32
33 if not current_platform.is_cpu() or not current_platform.is_zen_cpu():
34 pytest.skip("skipping non-Zen CPU tests", allow_module_level=True)
35
36 if not has_zentorch_op(["zentorch_fused_moe"]):
37 pytest.skip(
38 "skipping tests: zentorch_fused_moe op not available",
39 allow_module_level=True,
40 )
41
42 ZENTORCH_ACT = [
43 MoEActivation.SILU,
44 MoEActivation.GELU,
45 MoEActivation.GELU_TANH,
46 MoEActivation.SWIGLUOAI,
47 ]
48 USE_BIAS = [False, True]
49
50
51 def _make_layer_and_config(
52 act: MoEActivation,
53 *,
54 hidden_size: int = 128,
55 intermediate_size: int = 128,
56 expert_num: int = 8,
57 topk_num: int = 4,
58 use_bias: bool = False,
59 dtype: torch.dtype = torch.bfloat16,
60 ) -> tuple[_StubMoELayer, FusedMoEConfig]:
61 up_dim = 2 * intermediate_size
62 w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
63 0.5 * hidden_size**0.5
64 )
65 w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
66 0.5 * intermediate_size**0.5
67 )
68 w13_bias = None
69 w2_bias = None
70 if use_bias:
71 w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
72 w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
73 0.5 * hidden_size**0.5
74 )
75
76 layer = _StubMoELayer(
77 w13,
78 w2,
79 act,
80 w13_bias,
81 w2_bias,
82 )
83 moe_config = _make_moe_config(
84 expert_num,
85 hidden_size,
86 intermediate_size,
87 topk_num,
88 dtype,
89 act,
90 use_bias,
91 )
92 return layer, moe_config
93
94
95 @pytest.mark.parametrize("act", ZENTORCH_ACT)
96 def test_zen_cpu_fused_moe_config_supported(act: MoEActivation):
97 """CPUUnquantizedExperts is selected when zentorch MoE is supported."""
98 _, moe_config = _make_layer_and_config(act)
99 assert is_zentorch_moe_config_supported(moe_config)
100
101 supported, reason = CPUUnquantizedExperts.is_supported_config(
102 CPUUnquantizedExperts,
103 moe_config,
104 None,
105 None,
106 mk.FusedMoEActivationFormat.Standard,
107 )
108 assert supported, reason
109
110
111 def test_zen_cpu_fused_moe_config_supported_unaligned_intermediate():
112 """Zentorch accepts shapes the grouped-gemm kernel rejects."""
113 act = MoEActivation.SILU
114 moe_config = _make_moe_config(
115 expert_num=8,
116 hidden_size=128,
117 intermediate_size=UNALIGNED_INTERMEDIATE_DIM,
118 topk_num=4,
119 dtype=torch.bfloat16,
120 act=act,
121 )
122 assert is_zentorch_moe_config_supported(moe_config)
123
124 supported, reason = CPUUnquantizedExperts.is_supported_config(
125 CPUUnquantizedExperts,
126 moe_config,
127 None,
128 None,
129 mk.FusedMoEActivationFormat.Standard,
130 )
131 assert supported, reason
132
133
134 @pytest.mark.parametrize("act", ZENTORCH_ACT)
135 @pytest.mark.parametrize("use_bias", USE_BIAS)
136 def test_zen_cpu_fused_moe_forward(
137 default_vllm_config,
138 act: MoEActivation,
139 use_bias: bool,
140 ):
141 """The zentorch forward runs end-to-end and matches the reference."""
142 set_random_seed(0)
143
144 expert_num = 8
145 hidden_size = 128
146 batch_size = 64
147 dtype = torch.bfloat16
148 topk_num = expert_num // 2
149
150 layer, moe_config = _make_layer_and_config(
151 act,
152 expert_num=expert_num,
153 topk_num=topk_num,
154 use_bias=use_bias,
155 )
156
157 input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
158 0.5 * hidden_size**0.5
159 )
160 router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
161
162 quant_config = (
163 biased_moe_quant_config(layer.w13_bias, layer.w2_bias)
164 if use_bias
165 else FusedMoEQuantConfig.make()
166 )
167 experts = CPUUnquantizedExperts(moe_config, quant_config)
168 experts.process_weights_after_loading(layer)
169 assert experts._use_zentorch is True
170
171 w13_before = layer.w13_weight.detach().clone()
172 w2_before = layer.w2_weight.detach().clone()
173
174 topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False)
175
176 output = torch.empty_like(input)
177 experts.apply(
178 output=output,
179 hidden_states=input,
180 w1=layer.w13_weight,
181 w2=layer.w2_weight,
182 topk_weights=topk_weight,
183 topk_ids=topk_ids,
184 activation=act,
185 global_num_experts=expert_num,
186 expert_map=None,
187 a1q_scale=None,
188 a2_scale=None,
189 workspace13=torch.empty(0),
190 workspace2=torch.empty(0),
191 expert_tokens_meta=None,
192 apply_router_weight_on_input=False,
193 )
194
195 # zentorch needs no prepacking; weights are consumed as loaded.
196 torch.testing.assert_close(layer.w13_weight, w13_before)
197 torch.testing.assert_close(layer.w2_weight, w2_before)
198
199 ref_output = ref_fused_moe(
200 input,
201 layer.w13_weight,
202 layer.w2_weight,
203 getattr(layer, "w13_bias", None),
204 getattr(layer, "w2_bias", None),
205 topk_weight,
206 topk_ids,
207 act,
208 )
209
210 atol, rtol = get_default_atol(output), get_default_rtol(output)
211 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)