| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import sys |
| 5 | import time |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | |
| 10 | from vllm.platforms import CpuArchEnum, current_platform |
| 11 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 12 | from vllm.utils.torch_utils import set_random_seed |
| 13 | |
| 14 | # Check if CPU MoE operations are available |
| 15 | try: |
| 16 | from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight |
| 17 | except (ImportError, AttributeError) as e: |
| 18 | print(f"Import error: {e}") |
| 19 | sys.exit(1) |
| 20 | |
| 21 | # ISA selection following test_cpu_fused_moe.py pattern |
| 22 | ISA_CHOICES = ["vec"] |
| 23 | if torch.cpu._is_amx_tile_supported(): |
| 24 | ISA_CHOICES.append("amx") |
| 25 | if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: |
| 26 | ISA_CHOICES.append("neon") |
| 27 | |
| 28 | |
| 29 | @torch.inference_mode() |
| 30 | def main( |
| 31 | batch_size: int, |
| 32 | expert_num: int, |
| 33 | hidden_size: int, |
| 34 | intermediate_size: int, |
| 35 | topk_num: int, |
| 36 | use_bias: bool = False, |
| 37 | dtype: torch.dtype = torch.bfloat16, |
| 38 | activation: str = "silu", |
| 39 | isa: str = "vec", |
| 40 | seed: int = 0, |
| 41 | iters: int = 20, |
| 42 | ) -> None: |
| 43 | set_random_seed(seed) |
| 44 | # up_dim = 2 * intermediate_size for gate + up projection |
| 45 | up_dim = 2 * intermediate_size |
| 46 | |
| 47 | input_tensor = torch.randn((batch_size, hidden_size), dtype=dtype) / ( |
| 48 | 0.5 * hidden_size**0.5 |
| 49 | ) |
| 50 | |
| 51 | w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( |
| 52 | 0.5 * hidden_size**0.5 |
| 53 | ) |
| 54 | w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( |
| 55 | 0.5 * intermediate_size**0.5 |
| 56 | ) |
| 57 | |
| 58 | w13_bias = None |
| 59 | w2_bias = None |
| 60 | if use_bias: |
| 61 | w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) |
| 62 | w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( |
| 63 | 0.5 * hidden_size**0.5 |
| 64 | ) |
| 65 | |
| 66 | router_logits = torch.randn((batch_size, expert_num), dtype=dtype) |
| 67 | score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) |
| 68 | topk_weights, topk_ids = torch.topk(score, topk_num) |
| 69 | topk_ids = topk_ids.to(torch.int32) |
| 70 | |
| 71 | packed_w13 = cpu_prepack_moe_weight(w13, isa) |
| 72 | packed_w2 = cpu_prepack_moe_weight(w2, isa) |
| 73 | output = torch.empty_like(input_tensor) |
| 74 | |
| 75 | def run_benchmark(iters: int) -> list[float]: |
| 76 | times = [] |
| 77 | for _ in range(iters): |
| 78 | start_time = time.perf_counter_ns() |
| 79 | cpu_fused_moe( |
| 80 | output, |
| 81 | input_tensor, |
| 82 | packed_w13, |
| 83 | packed_w2, |
| 84 | w13_bias, |
| 85 | w2_bias, |
| 86 | topk_weights, |
| 87 | topk_ids, |
| 88 | activation, |
| 89 | isa, |
| 90 | ) |
| 91 | end_time = time.perf_counter_ns() |
| 92 | times.append((end_time - start_time) / 1e6) |
| 93 | return times |
| 94 | |
| 95 | # warmup |
| 96 | run_benchmark(5) |
| 97 | # benchmark |
| 98 | times = run_benchmark(iters) |
| 99 | |
| 100 | if not times: |
| 101 | print("No iterations to measure. Set --iters > 0.") |
| 102 | return |
| 103 | |
| 104 | time_min = min(times) |
| 105 | time_max = max(times) |
| 106 | time_mean = np.mean(times) |
| 107 | time_std = np.std(times) |
| 108 | |
| 109 | print("\tmin (ms) = ", time_min) |
| 110 | print("\tmax (ms) = ", time_max) |
| 111 | print("\tmean (ms) = ", time_mean) |
| 112 | print("\tstd = ", time_std) |
| 113 | print("\tmedian (ms) = ", np.median(times)) |
| 114 | |
| 115 | # Calculate throughput metrics |
| 116 | # FLOPs estimation: 2 * batch * topk * (hidden * up_dim + intermediate * hidden) |
| 117 | flops_per_token = ( |
| 118 | 2 * topk_num * (hidden_size * up_dim + intermediate_size * hidden_size) |
| 119 | ) |
| 120 | total_flops = batch_size * flops_per_token |
| 121 | tflops = total_flops / (time_mean * 1e-3) / 1e12 |
| 122 | print(f"\tthroughput (TFLOP/s) = {tflops:.4f}") |
| 123 | |
| 124 | |
| 125 | if __name__ == "__main__": |
| 126 | parser = FlexibleArgumentParser(description="Benchmark the CPU fused MoE kernel.") |
| 127 | parser.add_argument("--batch-size", type=int, default=64) |
| 128 | parser.add_argument("--expert-num", type=int, default=8) |
| 129 | parser.add_argument("--hidden-size", type=int, default=2880) |
| 130 | parser.add_argument("--intermediate-size", type=int, default=2880) |
| 131 | parser.add_argument( |
| 132 | "--topk-num", |
| 133 | type=int, |
| 134 | default=None, |
| 135 | help="Number of experts to route each token to (default: expert_num // 2)", |
| 136 | ) |
| 137 | parser.add_argument("--use-bias", action="store_true") |
| 138 | parser.add_argument( |
| 139 | "--activation", |
| 140 | type=str, |
| 141 | choices=["silu", "swigluoai"], |
| 142 | default="silu", |
| 143 | help="Activation function", |
| 144 | ) |
| 145 | parser.add_argument( |
| 146 | "--isa", |
| 147 | type=str, |
| 148 | choices=ISA_CHOICES, |
| 149 | default="vec", |
| 150 | help=f"ISA to use (available: {ISA_CHOICES})", |
| 151 | ) |
| 152 | parser.add_argument("--seed", type=int, default=0) |
| 153 | parser.add_argument("--iters", type=int, default=20) |
| 154 | |
| 155 | args = parser.parse_args() |
| 156 | |
| 157 | # Default topk_num to expert_num // 2, minimum 1 |
| 158 | topk_num = ( |
| 159 | args.topk_num if args.topk_num is not None else max(args.expert_num // 2, 1) |
| 160 | ) |
| 161 | |
| 162 | print(args) |
| 163 | |
| 164 | main( |
| 165 | batch_size=args.batch_size, |
| 166 | expert_num=args.expert_num, |
| 167 | hidden_size=args.hidden_size, |
| 168 | intermediate_size=args.intermediate_size, |
| 169 | topk_num=topk_num, |
| 170 | use_bias=args.use_bias, |
| 171 | dtype=torch.bfloat16, # Following test_cpu_fused_moe.py |
| 172 | activation=args.activation, |
| 173 | isa=args.isa, |
| 174 | seed=args.seed, |
| 175 | iters=args.iters, |
| 176 | ) |