@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 227 lines 6.78 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for CPU INT4 W4A8 dynamic quantized fused MoE kernel (CPUExpertsInt4)."""
4
5 import sys
6
7 import pytest
8 import torch
9 import torch.nn.functional as F
10
11 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12 from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
13 CPUExpertsInt4,
14 )
15 from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
16 convert_to_w4a8_int8_moe_format,
17 )
18 from vllm.platforms import CpuArchEnum, current_platform
19 from vllm.utils.torch_utils import set_random_seed
20
21 if (
22 not current_platform.is_cpu()
23 or current_platform.get_cpu_architecture() != CpuArchEnum.ARM
24 ):
25 pytest.skip("skipping Arm CPU-only tests", allow_module_level=True)
26
27
28 # Tolerance for INT4 W4A8
29 INT4_W4A8_ATOL = 2e-2
30 INT4_W4A8_RTOL = 2e-2
31
32
33 def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
34 """SwiGLU activation: SiLU(gate) * up."""
35 d = x.shape[-1] // 2
36 return F.silu(x[..., :d]) * x[..., d:]
37
38
39 def _make_int4_moe_weights(
40 E: int,
41 N: int,
42 K: int,
43 group_size: int,
44 has_bias: bool = False,
45 ) -> tuple[
46 torch.Tensor,
47 torch.Tensor,
48 torch.Tensor,
49 torch.Tensor,
50 torch.Tensor | None,
51 torch.Tensor | None,
52 ]:
53 """Generate random INT4 MoE weights with random scales.
54
55 Args:
56 E: Number of experts
57 N: Intermediate size
58 K: Hidden size
59 group_size: Quantization group size (-1 for channel-wise)
60 has_bias: Whether to include bias
61
62 Returns:
63 (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias)
64 where *_ref are the dequantized float reference weights
65
66 """
67 # Generate INT4 weights as int8 values in [-8, 7]
68 w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8)
69 w2_int4 = torch.randint(-8, 8, (E, K, N), dtype=torch.int8)
70
71 # Determine number of scale columns
72 def _n_scale_cols(in_features: int) -> int:
73 return 1 if group_size == -1 else (in_features // group_size)
74
75 # Generate random scales
76 scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
77 w13_scales = torch.rand(E, 2 * N, _n_scale_cols(K), dtype=scale_dtype) * 0.01
78 w2_scales = torch.rand(E, K, _n_scale_cols(N), dtype=scale_dtype) * 0.01
79
80 # Generate biases if needed
81 w13_bias = None
82 w2_bias = None
83 if has_bias:
84 w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01
85 w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01
86
87 w13_packed, w2_packed, *_ = convert_to_w4a8_int8_moe_format(
88 w13_weight=w13_int4,
89 w2_weight=w2_int4,
90 w13_weight_scale=w13_scales,
91 w2_weight_scale=w2_scales,
92 group_size=group_size,
93 w13_bias=w13_bias if has_bias else None,
94 w2_bias=w2_bias if has_bias else None,
95 )
96
97 if group_size == -1:
98 w13_scale = w13_scales.float()
99 w2_scale = w2_scales.float()
100 else:
101 w13_scale = w13_scales.float().repeat_interleave(group_size, dim=-1)
102 w2_scale = w2_scales.float().repeat_interleave(group_size, dim=-1)
103
104 w13_ref = w13_int4.float() * w13_scale
105 w2_ref = w2_int4.float() * w2_scale
106 if has_bias and w13_bias is not None:
107 w13_ref = w13_ref + w13_bias.float().unsqueeze(-1)
108 if has_bias and w2_bias is not None:
109 w2_ref = w2_ref + w2_bias.float().unsqueeze(-1)
110
111 return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias
112
113
114 def ref_int4_moe(
115 a: torch.Tensor,
116 w13_ref: torch.Tensor,
117 w2_ref: torch.Tensor,
118 topk_weight: torch.Tensor,
119 topk_ids: torch.Tensor,
120 ) -> torch.Tensor:
121 """Reference INT4 W4A8 fused MoE using dequantized weights.
122
123 Steps:
124 1. Use dequantized float weights
125 2. For each expert: matmul → SwiGLU → matmul
126 3. Weighted sum across top-k experts
127 """
128 B, D = a.shape
129 topk = topk_ids.size(1)
130
131 a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
132 out = torch.zeros(B * topk, w2_ref.shape[1], dtype=torch.float32)
133
134 topk_weight_flat = topk_weight.view(-1)
135 topk_ids_flat = topk_ids.view(-1)
136
137 for i in range(w13_ref.shape[0]):
138 mask = topk_ids_flat == i
139 if mask.sum():
140 # w13: [2N, K], input: [B, K] -> output: [B, 2N]
141 gate_up = torch.matmul(a_exp[mask], w13_ref[i].transpose(0, 1))
142 # SwiGLU activation
143 hidden = _silu_and_mul(gate_up)
144 # w2: [K, N], hidden: [B, N] -> output: [B, K]
145 out[mask] = torch.matmul(hidden, w2_ref[i].transpose(0, 1))
146
147 return (
148 (out.view(B, -1, w2_ref.shape[1]) * topk_weight_flat.view(B, -1, 1))
149 .sum(dim=1)
150 .to(a.dtype)
151 )
152
153
154 NUM_TOKENS = [1, 2, 64, 128]
155 # (intermediate_size N, hidden_size K, num_experts E, topk, group_size)
156 MoE_CONFIGS = [
157 (256, 512, 8, 2, 128),
158 (256, 512, 8, 2, 64),
159 (256, 512, 8, 2, -1), # channel-wise
160 (512, 256, 8, 4, 128),
161 (512, 512, 8, 2, 128),
162 (768, 2048, 8, 2, 128),
163 (768, 2048, 16, 4, 64),
164 ]
165 SEEDS = [0, 42]
166 ACTIVATION_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
167
168
169 @pytest.mark.parametrize("M", NUM_TOKENS)
170 @pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS)
171 @pytest.mark.parametrize("seed", SEEDS)
172 @pytest.mark.parametrize("activation_dtype", ACTIVATION_DTYPES)
173 def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed, activation_dtype):
174 """Test dynamic_4bit_int_moe kernel against dequantized torch reference."""
175 set_random_seed(seed)
176 activation = MoEActivation.SILU
177
178 # Generate input activations
179 a = torch.randn(M, K, dtype=activation_dtype) / (K**0.5)
180
181 # Generate INT4 weights
182 w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights(
183 E, N, K, group_size, has_bias=False
184 )
185
186 # Generate router logits and topk
187 score = torch.randn(M, E, dtype=torch.bfloat16)
188 score = torch.softmax(score, dim=-1, dtype=torch.float32)
189 topk_weight, topk_ids = torch.topk(score, topk)
190 topk_ids = topk_ids.to(torch.long)
191
192 # Reference output using dequantized weights
193 ref_out = ref_int4_moe(
194 a,
195 w13_ref,
196 w2_ref,
197 topk_weight,
198 topk_ids,
199 )
200
201 # Test dynamic_4bit_int_moe kernel
202 apply_router_weight_on_input = False
203
204 out = torch.ops._C.dynamic_4bit_int_moe(
205 a,
206 topk_ids,
207 topk_weight,
208 w13_packed,
209 w2_packed,
210 K, # H (hidden_size / w2_out_features)
211 N, # I (intermediate_size / w2_in_features)
212 group_size,
213 apply_router_weight_on_input,
214 CPUExpertsInt4._activation_kind(activation),
215 )
216
217 assert out.dtype == activation_dtype
218 torch.testing.assert_close(
219 ref_out,
220 out,
221 atol=INT4_W4A8_ATOL,
222 rtol=INT4_W4A8_RTOL,
223 )
224
225
226 if __name__ == "__main__":
227 sys.exit(pytest.main([__file__, "-v"]))