@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 252 lines 8.38 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the Zen CPU INT8 MoE experts (ZenCPUExpertsInt8)."""
4
5 import pytest
6 import torch
7
8 from tests.kernels.moe.test_cpu_fused_moe import (
9 _make_moe_config,
10 quantize_per_channel,
11 ref_fused_moe_int8,
12 )
13 from vllm.model_executor.kernels.linear.zentorch_utils import (
14 _ZENTORCH_MOE_ACTIVATIONS,
15 )
16 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17 from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEParallelConfig,
19 int8_w8a8_moe_quant_config,
20 )
21 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ZenCPUExpertsInt8
22 from vllm.model_executor.layers.fused_moe.oracle.int8 import (
23 Int8MoeBackend,
24 backend_to_kernel_cls,
25 )
26 from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk
27 from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 kInt8DynamicTokenSym,
29 kInt8StaticChannelSym,
30 )
31 from vllm.platforms import current_platform
32 from vllm.utils.torch_utils import set_random_seed
33
34 if not current_platform.is_cpu():
35 pytest.skip("skipping CPU-only tests", allow_module_level=True)
36
37 EXPERT_NUM = 8
38 HIDDEN_SIZE = 128
39 INTERMEDIATE_SIZE = 128
40 BATCH_SIZE = 32
41 # zentorch names its activations as strings; the experts API is enum-typed.
42 ZENTORCH_ACT = [act for act in MoEActivation if act.value in _ZENTORCH_MOE_ACTIVATIONS]
43
44
45 @pytest.fixture
46 def mock_zentorch_fused_moe():
47 """Reference zentorch_fused_moe, registered when zentorch is absent."""
48 if hasattr(torch.ops.zentorch, "zentorch_fused_moe"):
49 yield None
50 return
51
52 calls: dict[str, dict] = {}
53
54 def _fused_moe(
55 output,
56 input,
57 w13,
58 w2,
59 w13_bias,
60 w2_bias,
61 topk_weights,
62 topk_id,
63 skip_weighted,
64 act,
65 w13_scales=None,
66 w2_scales=None,
67 zentorch_op_name="",
68 ):
69 calls["fused_moe"] = {
70 "w13": w13,
71 "w2": w2,
72 "w13_bias": w13_bias,
73 "w2_bias": w2_bias,
74 "topk_weights": topk_weights,
75 "topk_id": topk_id,
76 "skip_weighted": skip_weighted,
77 "act": act,
78 "w13_scales": w13_scales,
79 "w2_scales": w2_scales,
80 }
81 output.copy_(
82 ref_fused_moe_int8(
83 input,
84 w13,
85 w2,
86 w13_scales.float(),
87 w2_scales.float(),
88 None if w13_bias is None else w13_bias.float(),
89 None if w2_bias is None else w2_bias.float(),
90 topk_weights,
91 topk_id,
92 MoEActivation(act),
93 )
94 )
95
96 lib_def = torch.library.Library("zentorch", "DEF")
97 lib_def.define(
98 "zentorch_fused_moe(Tensor(a!) output, Tensor input, Tensor w13, "
99 "Tensor w2, Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, "
100 "Tensor topk_id, bool skip_weighted, str act, "
101 "Tensor? w13_scales=None, Tensor? w2_scales=None, *, "
102 "str zentorch_op_name='zentorch::zentorch_fused_moe') -> ()"
103 )
104 lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
105 lib_impl.impl("zentorch_fused_moe", _fused_moe)
106
107 yield calls
108
109 lib_impl._destroy()
110 lib_def._destroy()
111
112
113 def _make_int8_moe_weights(use_bias: bool, dtype: torch.dtype = torch.bfloat16):
114 up_dim = 2 * INTERMEDIATE_SIZE
115 w13_fp = torch.randn((EXPERT_NUM, up_dim, HIDDEN_SIZE), dtype=dtype) / (
116 0.5 * HIDDEN_SIZE**0.5
117 )
118 w2_fp = torch.randn((EXPERT_NUM, HIDDEN_SIZE, INTERMEDIATE_SIZE), dtype=dtype) / (
119 0.5 * INTERMEDIATE_SIZE**0.5
120 )
121 w13, w13_scale = quantize_per_channel(w13_fp)
122 w2, w2_scale = quantize_per_channel(w2_fp)
123 w13_scale = w13_scale.to(torch.bfloat16).float()
124 w2_scale = w2_scale.to(torch.bfloat16).float()
125
126 w13_bias = w2_bias = None
127 if use_bias:
128 w13_bias = torch.randn((EXPERT_NUM, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
129 w2_bias = torch.randn((EXPERT_NUM, HIDDEN_SIZE), dtype=dtype) / (
130 0.5 * HIDDEN_SIZE**0.5
131 )
132 return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias
133
134
135 def test_zen_int8_first_in_cpu_backend_priority():
136 """The Zen experts are preferred over the generic CPU int8 experts."""
137 assert backend_to_kernel_cls(Int8MoeBackend.CPU)[0] is ZenCPUExpertsInt8
138
139
140 def test_zen_int8_support_predicates(mock_zentorch_fused_moe, monkeypatch):
141 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: False)
142 assert not ZenCPUExpertsInt8._supports_current_device()
143 monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
144 assert ZenCPUExpertsInt8._supports_current_device()
145 assert ZenCPUExpertsInt8._supports_quant_scheme(
146 kInt8StaticChannelSym, kInt8DynamicTokenSym
147 )
148 assert not ZenCPUExpertsInt8._supports_quant_scheme(kInt8StaticChannelSym, None)
149 for act in ZENTORCH_ACT:
150 assert ZenCPUExpertsInt8._supports_activation(act)
151 assert not ZenCPUExpertsInt8._supports_activation(MoEActivation.RELU2)
152 assert not ZenCPUExpertsInt8._supports_activation(
153 MoEActivation.SWIGLUOAI_UNINTERLEAVE
154 )
155 assert not ZenCPUExpertsInt8._supports_no_act_and_mul()
156 assert ZenCPUExpertsInt8._supports_parallel_config(
157 FusedMoEParallelConfig.make_no_parallel()
158 )
159
160
161 @pytest.mark.parametrize("act", ZENTORCH_ACT)
162 @pytest.mark.parametrize("use_bias", [False, True])
163 def test_zen_int8_dispatch_contract(
164 default_vllm_config,
165 mock_zentorch_fused_moe,
166 act: MoEActivation,
167 use_bias: bool,
168 ):
169 """apply() hands zentorch the weights, scales and biases it expects."""
170 if mock_zentorch_fused_moe is None:
171 pytest.skip("real zentorch is installed; the recorded call is unavailable")
172 set_random_seed(0)
173
174 topk_num = EXPERT_NUM // 2
175 w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = _make_int8_moe_weights(use_bias)
176 moe_config = _make_moe_config(
177 EXPERT_NUM,
178 HIDDEN_SIZE,
179 INTERMEDIATE_SIZE,
180 topk_num,
181 torch.bfloat16,
182 act,
183 use_bias,
184 )
185 # A layer carries per-channel scales as [E, N, 1].
186 quant_config = int8_w8a8_moe_quant_config(
187 w1_scale=w13_scale.unsqueeze(-1),
188 w2_scale=w2_scale.unsqueeze(-1),
189 a1_scale=None,
190 a2_scale=None,
191 w1_bias=w13_bias,
192 w2_bias=w2_bias,
193 per_act_token_quant=True,
194 )
195
196 input = torch.randn((BATCH_SIZE, HIDDEN_SIZE), dtype=torch.bfloat16) / (
197 0.5 * HIDDEN_SIZE**0.5
198 )
199 router_logits = torch.randn((BATCH_SIZE, EXPERT_NUM), dtype=torch.bfloat16)
200 topk_weights, topk_ids = _softmax_topk(router_logits, topk_num, False)
201
202 experts = ZenCPUExpertsInt8(moe_config, quant_config)
203 # zentorch_fused_moe takes no expert_map.
204 assert not experts.supports_expert_map()
205 output = torch.empty_like(input)
206 experts.apply(
207 output=output,
208 hidden_states=input,
209 w1=w13,
210 w2=w2,
211 topk_weights=topk_weights,
212 topk_ids=topk_ids,
213 activation=act,
214 global_num_experts=EXPERT_NUM,
215 expert_map=None,
216 a1q_scale=None,
217 a2_scale=None,
218 workspace13=torch.empty(0),
219 workspace2=torch.empty(0),
220 expert_tokens_meta=None,
221 apply_router_weight_on_input=False,
222 )
223
224 call = mock_zentorch_fused_moe["fused_moe"]
225 assert call["w13_scales"].shape == (EXPERT_NUM, 2 * INTERMEDIATE_SIZE)
226 assert call["w2_scales"].shape == (EXPERT_NUM, HIDDEN_SIZE)
227 # f32 scales keep the layer on zentorch's fast routed MoE executor.
228 assert call["w13_scales"].dtype == torch.float32
229 assert call["w2_scales"].dtype == torch.float32
230 assert call["act"] == act.value.lower()
231 assert call["skip_weighted"] is False
232 assert call["topk_weights"].dtype == torch.float32
233 assert call["topk_id"].dtype == torch.int32
234 if use_bias:
235 assert call["w13_bias"].shape == (EXPERT_NUM, 2 * INTERMEDIATE_SIZE)
236 assert call["w2_bias"].shape == (EXPERT_NUM, HIDDEN_SIZE)
237 else:
238 assert call["w13_bias"] is None and call["w2_bias"] is None
239
240 ref = ref_fused_moe_int8(
241 input,
242 w13,
243 w2,
244 w13_scale,
245 w2_scale,
246 None if w13_bias is None else w13_bias.float(),
247 None if w2_bias is None else w2_bias.float(),
248 topk_weights.to(torch.float32),
249 topk_ids.to(torch.int32),
250 act,
251 )
252 torch.testing.assert_close(output, ref)