| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 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 vllm._custom_ops import ( |
| 11 | cpu_fused_moe, |
| 12 | cpu_fused_moe_int8, |
| 13 | cpu_prepack_moe_weight, |
| 14 | cpu_prepack_moe_weight_int8, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 17 | from vllm.model_executor.layers.fused_moe.config import ( |
| 18 | FusedMoEConfig, |
| 19 | FusedMoEParallelConfig, |
| 20 | FusedMoEQuantConfig, |
| 21 | RoutingMethodType, |
| 22 | biased_moe_quant_config, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 25 | ArmCPUUnquantizedExperts, |
| 26 | CPUUnquantizedExperts, |
| 27 | PowerCPUUnquantizedExperts, |
| 28 | X86CPUUnquantizedExperts, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk |
| 31 | from vllm.platforms import CpuArchEnum, 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 = [ |
| 38 | 8, |
| 39 | ] |
| 40 | HIDDEN_DIM = [128, 2880] |
| 41 | INTERMEDIATE_DIM = [128, 2880] |
| 42 | BATCH_SIZE = [1, 64, 256] |
| 43 | ACT = [ |
| 44 | MoEActivation.SILU, |
| 45 | MoEActivation.SWIGLUOAI, |
| 46 | MoEActivation.GELU, |
| 47 | MoEActivation.GELU_TANH, |
| 48 | ] |
| 49 | USE_BIAS = [False, True] |
| 50 | ISA = ["vec"] |
| 51 | if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: |
| 52 | ISA.append("neon") |
| 53 | if torch.cpu._is_amx_tile_supported(): |
| 54 | ISA.append("amx") |
| 55 | if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: |
| 56 | ISA.append("vsx") |
| 57 | |
| 58 | DTYPE = [torch.bfloat16] |
| 59 | |
| 60 | |
| 61 | def _ref_moe_activation( |
| 62 | input: torch.Tensor, |
| 63 | activation: MoEActivation, |
| 64 | ) -> torch.Tensor: |
| 65 | if activation == MoEActivation.SWIGLUOAI: |
| 66 | gate = input[..., ::2].clamp(max=7.0) |
| 67 | up = input[..., 1::2].clamp(min=-7.0, max=7.0) |
| 68 | glu = gate * torch.sigmoid(gate * 1.702) |
| 69 | return (up + 1) * glu |
| 70 | |
| 71 | d = input.shape[-1] // 2 |
| 72 | gate, up = input[..., :d], input[..., d:] |
| 73 | if activation == MoEActivation.SILU: |
| 74 | return torch.nn.functional.silu(gate) * up |
| 75 | if activation == MoEActivation.GELU: |
| 76 | return torch.nn.functional.gelu(gate, approximate="none") * up |
| 77 | if activation == MoEActivation.GELU_TANH: |
| 78 | return torch.nn.functional.gelu(gate, approximate="tanh") * up |
| 79 | raise ValueError(f"Unsupported activation: {activation}") |
| 80 | |
| 81 | |
| 82 | def ref_fused_moe( |
| 83 | input: torch.Tensor, |
| 84 | w13: torch.Tensor, |
| 85 | w2: torch.Tensor, |
| 86 | w13_bias: torch.Tensor | None, |
| 87 | w2_bias: torch.Tensor | None, |
| 88 | topk_weights: torch.Tensor, |
| 89 | topk_ids: torch.Tensor, |
| 90 | activation: MoEActivation, |
| 91 | ) -> torch.Tensor: |
| 92 | len_experts = w13.size(0) |
| 93 | |
| 94 | cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts)) |
| 95 | cnts.scatter_(1, topk_ids.to(torch.int64), 1) |
| 96 | tokens_per_expert = cnts.sum(dim=0) |
| 97 | idxs = topk_ids.view(-1).argsort() |
| 98 | |
| 99 | sorted_tokens = input[idxs // topk_ids.shape[1]] |
| 100 | tokens_per_expert = tokens_per_expert.cpu().numpy() |
| 101 | |
| 102 | outputs = [] |
| 103 | start_idx = 0 |
| 104 | |
| 105 | for i, num_tokens in enumerate(tokens_per_expert): |
| 106 | end_idx = start_idx + num_tokens |
| 107 | if num_tokens == 0: |
| 108 | continue |
| 109 | tokens_for_this_expert = sorted_tokens[start_idx:end_idx].float() |
| 110 | curr_w13 = w13[i].float() |
| 111 | curr_w2 = w2[i].float() |
| 112 | |
| 113 | curr_w13_bias = None |
| 114 | if w13_bias is not None: |
| 115 | curr_w13_bias = w13_bias[i].float() |
| 116 | |
| 117 | curr_w2_bias = None |
| 118 | if w2_bias is not None: |
| 119 | curr_w2_bias = w2_bias[i].float() |
| 120 | |
| 121 | gate_up = torch.nn.functional.linear( |
| 122 | tokens_for_this_expert, curr_w13, curr_w13_bias |
| 123 | ) |
| 124 | # Note: to simulate the kernel implementation |
| 125 | gate_up = _ref_moe_activation(gate_up, activation) |
| 126 | gate_up = gate_up.to(dtype=input.dtype).float() |
| 127 | expert_out = torch.nn.functional.linear(gate_up, curr_w2, curr_w2_bias) |
| 128 | |
| 129 | outputs.append(expert_out) |
| 130 | start_idx = end_idx |
| 131 | |
| 132 | outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) |
| 133 | new_x = torch.empty_like(outs) |
| 134 | |
| 135 | new_x[idxs] = outs |
| 136 | final_out = ( |
| 137 | new_x.view(*topk_ids.shape, -1) |
| 138 | .mul_(topk_weights.unsqueeze(dim=-1)) |
| 139 | .sum(dim=1) |
| 140 | .type(input.dtype) |
| 141 | ) |
| 142 | return final_out |
| 143 | |
| 144 | |
| 145 | def quantize_per_channel( |
| 146 | weight: torch.Tensor, |
| 147 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 148 | """Symmetrically quantize each weight output channel.""" |
| 149 | weight_fp32 = weight.float() |
| 150 | scale = weight_fp32.abs().amax(dim=-1).clamp_min(1e-12) / 127.0 |
| 151 | quantized = ( |
| 152 | (weight_fp32 / scale.unsqueeze(-1)).round().clamp(-127, 127).to(torch.int8) |
| 153 | ) |
| 154 | return quantized, scale |
| 155 | |
| 156 | |
| 157 | def quantize_per_token( |
| 158 | input: torch.Tensor, |
| 159 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 160 | """Symmetrically quantize each input token.""" |
| 161 | input_fp32 = input.float() |
| 162 | scale = input_fp32.abs().amax(dim=-1, keepdim=True).clamp_min(1e-7) / 127.0 |
| 163 | quantized = (input_fp32 / scale).round().clamp(-127, 127).to(torch.int8) |
| 164 | return quantized, scale |
| 165 | |
| 166 | |
| 167 | def ref_fused_moe_int8( |
| 168 | input: torch.Tensor, |
| 169 | w13: torch.Tensor, |
| 170 | w2: torch.Tensor, |
| 171 | w13_scale: torch.Tensor, |
| 172 | w2_scale: torch.Tensor, |
| 173 | w13_bias: torch.Tensor | None, |
| 174 | w2_bias: torch.Tensor | None, |
| 175 | topk_weights: torch.Tensor, |
| 176 | topk_ids: torch.Tensor, |
| 177 | activation: MoEActivation, |
| 178 | ) -> torch.Tensor: |
| 179 | """Reference the two dynamically quantized INT8 GEMMs.""" |
| 180 | input_int8, input_scale = quantize_per_token(input) |
| 181 | expert_num = w13.size(0) |
| 182 | |
| 183 | counts = topk_ids.new_zeros((topk_ids.size(0), expert_num)) |
| 184 | counts.scatter_(1, topk_ids.to(torch.int64), 1) |
| 185 | tokens_per_expert = counts.sum(dim=0).tolist() |
| 186 | sorted_route_ids = topk_ids.view(-1).argsort() |
| 187 | sorted_token_ids = sorted_route_ids // topk_ids.size(1) |
| 188 | |
| 189 | outputs = [] |
| 190 | start_idx = 0 |
| 191 | for expert_idx, token_count in enumerate(tokens_per_expert): |
| 192 | end_idx = start_idx + token_count |
| 193 | if token_count == 0: |
| 194 | continue |
| 195 | |
| 196 | token_ids = sorted_token_ids[start_idx:end_idx] |
| 197 | gate_up = torch.matmul( |
| 198 | input_int8[token_ids].float(), |
| 199 | w13[expert_idx].float().T, |
| 200 | ) |
| 201 | gate_up *= input_scale[token_ids] * w13_scale[expert_idx] |
| 202 | if w13_bias is not None: |
| 203 | gate_up += w13_bias[expert_idx].float() |
| 204 | |
| 205 | intermediate = _ref_moe_activation(gate_up, activation).to(input.dtype) |
| 206 | intermediate_int8, intermediate_scale = quantize_per_token(intermediate) |
| 207 | output = torch.matmul( |
| 208 | intermediate_int8.float(), |
| 209 | w2[expert_idx].float().T, |
| 210 | ) |
| 211 | output *= intermediate_scale * w2_scale[expert_idx] |
| 212 | if w2_bias is not None: |
| 213 | output += w2_bias[expert_idx].float() |
| 214 | |
| 215 | outputs.append(output) |
| 216 | start_idx = end_idx |
| 217 | |
| 218 | sorted_output = torch.cat(outputs, dim=0) |
| 219 | routed_output = torch.empty_like(sorted_output) |
| 220 | routed_output[sorted_route_ids] = sorted_output |
| 221 | return ( |
| 222 | routed_output.view(*topk_ids.shape, -1) |
| 223 | .mul_(topk_weights.unsqueeze(-1)) |
| 224 | .sum(dim=1) |
| 225 | .to(input.dtype) |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | @pytest.mark.parametrize("batch_size", BATCH_SIZE) |
| 230 | @pytest.mark.parametrize("expert_num", EXPERT_NUM) |
| 231 | @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) |
| 232 | @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) |
| 233 | @pytest.mark.parametrize("use_bias", USE_BIAS) |
| 234 | @pytest.mark.parametrize("dtype", DTYPE) |
| 235 | @pytest.mark.parametrize("act", ACT) |
| 236 | @pytest.mark.parametrize("isa", ISA) |
| 237 | def test_cpu_fused_moe( |
| 238 | default_vllm_config, |
| 239 | batch_size: int, |
| 240 | expert_num: int, |
| 241 | hidden_size: int, |
| 242 | intermediate_size: int, |
| 243 | use_bias: bool, |
| 244 | dtype: torch.dtype, |
| 245 | act: MoEActivation, |
| 246 | isa: str, |
| 247 | ): |
| 248 | set_random_seed(0) |
| 249 | |
| 250 | topk_num = max(expert_num // 2, 1) |
| 251 | up_dim = 2 * intermediate_size |
| 252 | |
| 253 | input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( |
| 254 | 0.5 * hidden_size**0.5 |
| 255 | ) |
| 256 | w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( |
| 257 | 0.5 * hidden_size**0.5 |
| 258 | ) |
| 259 | w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( |
| 260 | 0.5 * intermediate_size**0.5 |
| 261 | ) |
| 262 | router_logits = torch.randn((batch_size, expert_num), dtype=dtype) |
| 263 | w13_bias = None |
| 264 | w2_bias = None |
| 265 | if use_bias: |
| 266 | w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) |
| 267 | w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( |
| 268 | 0.5 * hidden_size**0.5 |
| 269 | ) |
| 270 | score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 271 | topk_weight, topk_ids = torch.topk(score, topk_num) |
| 272 | topk_ids = topk_ids.to(torch.int32) |
| 273 | |
| 274 | ref_output = ref_fused_moe( |
| 275 | input, |
| 276 | w13, |
| 277 | w2, |
| 278 | w13_bias, |
| 279 | w2_bias, |
| 280 | topk_weight, |
| 281 | topk_ids, |
| 282 | act, |
| 283 | ) |
| 284 | |
| 285 | packed_w13 = cpu_prepack_moe_weight(w13, isa) |
| 286 | packed_w2 = cpu_prepack_moe_weight(w2, isa) |
| 287 | output = torch.empty_like(input) |
| 288 | cpu_fused_moe( |
| 289 | output, |
| 290 | input, |
| 291 | packed_w13, |
| 292 | packed_w2, |
| 293 | w13_bias, |
| 294 | w2_bias, |
| 295 | topk_weight, |
| 296 | topk_ids, |
| 297 | act.value, |
| 298 | isa, |
| 299 | ) |
| 300 | |
| 301 | atol, rtol = get_default_atol(output), get_default_rtol(output) |
| 302 | ( |
| 303 | torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), |
| 304 | f"{torch.max(torch.abs(output - ref_output))}", |
| 305 | ) |
| 306 | |
| 307 | |
| 308 | @pytest.mark.skipif( |
| 309 | current_platform.get_cpu_architecture() |
| 310 | not in (CpuArchEnum.ARM, CpuArchEnum.POWERPC), |
| 311 | reason="Requires Arm or POWER CPU", |
| 312 | ) |
| 313 | @pytest.mark.parametrize("batch_size", BATCH_SIZE) |
| 314 | @pytest.mark.parametrize("expert_num", EXPERT_NUM) |
| 315 | @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) |
| 316 | @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) |
| 317 | @pytest.mark.parametrize("use_bias", USE_BIAS) |
| 318 | @pytest.mark.parametrize( |
| 319 | "dtype", |
| 320 | ( |
| 321 | [torch.float32, torch.float16, torch.bfloat16] |
| 322 | if current_platform.get_cpu_architecture() == CpuArchEnum.ARM |
| 323 | else [torch.float32, torch.bfloat16] |
| 324 | ), |
| 325 | ) |
| 326 | @pytest.mark.parametrize("act", ACT) |
| 327 | @pytest.mark.parametrize( |
| 328 | "isa", |
| 329 | ["neon"] if current_platform.get_cpu_architecture() == CpuArchEnum.ARM else ["vsx"], |
| 330 | ) |
| 331 | def test_cpu_fused_moe_int8( |
| 332 | batch_size: int, |
| 333 | expert_num: int, |
| 334 | hidden_size: int, |
| 335 | intermediate_size: int, |
| 336 | use_bias: bool, |
| 337 | dtype: torch.dtype, |
| 338 | act: MoEActivation, |
| 339 | isa: str, |
| 340 | ): |
| 341 | set_random_seed(0) |
| 342 | topk_num = max(expert_num // 2, 1) |
| 343 | up_dim = 2 * intermediate_size |
| 344 | |
| 345 | input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( |
| 346 | 0.5 * hidden_size**0.5 |
| 347 | ) |
| 348 | w13_fp = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( |
| 349 | 0.5 * hidden_size**0.5 |
| 350 | ) |
| 351 | w2_fp = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( |
| 352 | 0.5 * intermediate_size**0.5 |
| 353 | ) |
| 354 | w13, w13_scale = quantize_per_channel(w13_fp) |
| 355 | w2, w2_scale = quantize_per_channel(w2_fp) |
| 356 | |
| 357 | w13_bias = None |
| 358 | w2_bias = None |
| 359 | if use_bias: |
| 360 | w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) |
| 361 | w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( |
| 362 | 0.5 * hidden_size**0.5 |
| 363 | ) |
| 364 | |
| 365 | router_logits = torch.randn((batch_size, expert_num), dtype=dtype) |
| 366 | score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 367 | topk_weights, topk_ids = torch.topk(score, topk_num) |
| 368 | topk_ids = topk_ids.to(torch.int32) |
| 369 | |
| 370 | ref_output = ref_fused_moe_int8( |
| 371 | input, |
| 372 | w13, |
| 373 | w2, |
| 374 | w13_scale, |
| 375 | w2_scale, |
| 376 | w13_bias, |
| 377 | w2_bias, |
| 378 | topk_weights, |
| 379 | topk_ids, |
| 380 | act, |
| 381 | ) |
| 382 | packed_w13 = cpu_prepack_moe_weight_int8(w13, isa) |
| 383 | packed_w2 = cpu_prepack_moe_weight_int8(w2, isa) |
| 384 | output = torch.empty_like(input) |
| 385 | cpu_fused_moe_int8( |
| 386 | output, |
| 387 | input, |
| 388 | packed_w13, |
| 389 | packed_w2, |
| 390 | w13_scale, |
| 391 | w2_scale, |
| 392 | w13_bias, |
| 393 | w2_bias, |
| 394 | topk_weights, |
| 395 | topk_ids, |
| 396 | act.value, |
| 397 | isa, |
| 398 | ) |
| 399 | |
| 400 | atol = rtol = 1e-1 if isa == "vsx" else 2e-2 |
| 401 | torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) |
| 402 | |
| 403 | |
| 404 | # moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel |
| 405 | # sharding of moe_intermediate_size=704 at tp=4 produces (704 // 4 == 176). |
| 406 | UNALIGNED_INTERMEDIATE_DIM = 176 |
| 407 | |
| 408 | |
| 409 | class _StubMoELayer(torch.nn.Module): |
| 410 | """Minimal stand-in for the real MoE layer module, exposing just what |
| 411 | the unquantized CPU experts read/replace (w13_weight, w2_weight, the |
| 412 | router configuration, and optionally w13_bias/w2_bias).""" |
| 413 | |
| 414 | def __init__( |
| 415 | self, |
| 416 | w13_weight: torch.Tensor, |
| 417 | w2_weight: torch.Tensor, |
| 418 | activation: MoEActivation, |
| 419 | w13_bias: torch.Tensor | None = None, |
| 420 | w2_bias: torch.Tensor | None = None, |
| 421 | ): |
| 422 | super().__init__() |
| 423 | self.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) |
| 424 | self.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) |
| 425 | self.activation = activation |
| 426 | if w13_bias is not None: |
| 427 | self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) |
| 428 | if w2_bias is not None: |
| 429 | self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) |
| 430 | self.use_grouped_topk = False |
| 431 | self.renormalize = False |
| 432 | self.scoring_func = "softmax" |
| 433 | self.custom_routing_function = None |
| 434 | |
| 435 | |
| 436 | def _make_moe_config( |
| 437 | expert_num: int, |
| 438 | hidden_size: int, |
| 439 | intermediate_size: int, |
| 440 | topk_num: int, |
| 441 | dtype: torch.dtype, |
| 442 | act: MoEActivation, |
| 443 | has_bias: bool = False, |
| 444 | ) -> FusedMoEConfig: |
| 445 | return FusedMoEConfig( |
| 446 | num_experts=expert_num, |
| 447 | experts_per_token=topk_num, |
| 448 | hidden_dim=hidden_size, |
| 449 | intermediate_size=intermediate_size, |
| 450 | num_local_experts=expert_num, |
| 451 | num_logical_experts=expert_num, |
| 452 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 453 | activation=act, |
| 454 | in_dtype=dtype, |
| 455 | device="cpu", |
| 456 | routing_method=RoutingMethodType.Default, |
| 457 | has_bias=has_bias, |
| 458 | ) |
| 459 | |
| 460 | |
| 461 | @pytest.mark.parametrize( |
| 462 | ("hidden_size", "intermediate_size", "act", "expected"), |
| 463 | [ |
| 464 | (128, 128, MoEActivation.SILU, True), |
| 465 | # Hidden size is not 32-aligned. |
| 466 | (112, 128, MoEActivation.SILU, False), |
| 467 | (128, 176, MoEActivation.SILU, True), |
| 468 | # SwiGLUOAI's interleaved gate/up layout cannot be padded. |
| 469 | (128, 176, MoEActivation.SWIGLUOAI, False), |
| 470 | ], |
| 471 | ) |
| 472 | def test_cpu_vec_fused_moe_shape_support( |
| 473 | hidden_size: int, |
| 474 | intermediate_size: int, |
| 475 | act: MoEActivation, |
| 476 | expected: bool, |
| 477 | ): |
| 478 | moe_config = _make_moe_config( |
| 479 | expert_num=8, |
| 480 | hidden_size=hidden_size, |
| 481 | intermediate_size=intermediate_size, |
| 482 | topk_num=4, |
| 483 | dtype=torch.bfloat16, |
| 484 | act=act, |
| 485 | ) |
| 486 | supported, _ = CPUUnquantizedExperts.is_supported_config( |
| 487 | CPUUnquantizedExperts, |
| 488 | moe_config, |
| 489 | None, |
| 490 | None, |
| 491 | mk.FusedMoEActivationFormat.Standard, |
| 492 | ) |
| 493 | assert supported is expected |
| 494 | |
| 495 | |
| 496 | @pytest.mark.parametrize("expert_num", EXPERT_NUM) |
| 497 | @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) |
| 498 | @pytest.mark.parametrize("use_bias", USE_BIAS) |
| 499 | @pytest.mark.parametrize("dtype", DTYPE) |
| 500 | @pytest.mark.parametrize("isa", ISA) |
| 501 | @pytest.mark.parametrize( |
| 502 | "act", |
| 503 | [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH], |
| 504 | ) |
| 505 | def test_cpu_fused_moe_unaligned_intermediate_size( |
| 506 | default_vllm_config, |
| 507 | expert_num: int, |
| 508 | hidden_size: int, |
| 509 | use_bias: bool, |
| 510 | dtype: torch.dtype, |
| 511 | isa: str, |
| 512 | act: MoEActivation, |
| 513 | ): |
| 514 | """CPU kernels handle unaligned intermediate sizes by zero-padding the |
| 515 | weights before prepacking.""" |
| 516 | set_random_seed(0) |
| 517 | batch_size = 64 |
| 518 | intermediate_size = UNALIGNED_INTERMEDIATE_DIM |
| 519 | topk_num = max(expert_num // 2, 1) |
| 520 | up_dim = 2 * intermediate_size |
| 521 | |
| 522 | input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( |
| 523 | 0.5 * hidden_size**0.5 |
| 524 | ) |
| 525 | w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( |
| 526 | 0.5 * hidden_size**0.5 |
| 527 | ) |
| 528 | w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( |
| 529 | 0.5 * intermediate_size**0.5 |
| 530 | ) |
| 531 | router_logits = torch.randn((batch_size, expert_num), dtype=dtype) |
| 532 | w13_bias = None |
| 533 | w2_bias = None |
| 534 | if use_bias: |
| 535 | w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) |
| 536 | w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( |
| 537 | 0.5 * hidden_size**0.5 |
| 538 | ) |
| 539 | # Route with the same helper the router uses internally, so the |
| 540 | # reference only differs from the kernel in how the experts are |
| 541 | # evaluated. |
| 542 | topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False) |
| 543 | |
| 544 | ref_output = ref_fused_moe( |
| 545 | input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act |
| 546 | ) |
| 547 | |
| 548 | layer = _StubMoELayer( |
| 549 | w13.clone(), |
| 550 | w2.clone(), |
| 551 | act, |
| 552 | w13_bias.clone() if w13_bias is not None else None, |
| 553 | w2_bias.clone() if w2_bias is not None else None, |
| 554 | ) |
| 555 | |
| 556 | moe_config = _make_moe_config( |
| 557 | expert_num, hidden_size, intermediate_size, topk_num, dtype, act, use_bias |
| 558 | ) |
| 559 | experts_cls = { |
| 560 | "vec": CPUUnquantizedExperts, |
| 561 | "amx": X86CPUUnquantizedExperts, |
| 562 | "neon": ArmCPUUnquantizedExperts, |
| 563 | "vsx": PowerCPUUnquantizedExperts, |
| 564 | }[isa] |
| 565 | supported, reason = experts_cls.is_supported_config( |
| 566 | experts_cls, |
| 567 | moe_config, |
| 568 | None, |
| 569 | None, |
| 570 | mk.FusedMoEActivationFormat.Standard, |
| 571 | ) |
| 572 | assert supported, ( |
| 573 | f"expected the padded intermediate size to be supported, got {reason}" |
| 574 | ) |
| 575 | |
| 576 | # Mirror UnquantizedFusedMoEMethod._setup_kernel: the quant config is |
| 577 | # built from the layer first, then the experts shuffle the layer into the |
| 578 | # kernel's runtime format. |
| 579 | quant_config = ( |
| 580 | biased_moe_quant_config(layer.w13_bias, layer.w2_bias) |
| 581 | if use_bias |
| 582 | else FusedMoEQuantConfig.make() |
| 583 | ) |
| 584 | experts = experts_cls(moe_config, quant_config) |
| 585 | assert experts.isa == isa |
| 586 | experts.process_weights_after_loading(layer) |
| 587 | |
| 588 | output = torch.empty_like(input) |
| 589 | experts.apply( |
| 590 | output=output, |
| 591 | hidden_states=input, |
| 592 | w1=layer.w13_weight, |
| 593 | w2=layer.w2_weight, |
| 594 | topk_weights=topk_weight, |
| 595 | topk_ids=topk_ids, |
| 596 | activation=act, |
| 597 | global_num_experts=expert_num, |
| 598 | expert_map=None, |
| 599 | a1q_scale=None, |
| 600 | a2_scale=None, |
| 601 | workspace13=torch.empty(0), |
| 602 | workspace2=torch.empty(0), |
| 603 | expert_tokens_meta=None, |
| 604 | apply_router_weight_on_input=False, |
| 605 | ) |
| 606 | |
| 607 | atol, rtol = get_default_atol(output), get_default_rtol(output) |
| 608 | torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) |