A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 07:45 UTC
[cpu][bench] Add Fused MoE Micro Benchmark for CPU Backend (#32092)
Signed-off-by: andikarachman <andika.rachman.y@gmail.com>
Andika Rachman committed
Jan 12, 2026 at 17:03 UTC
5e034f2e3d6a90b6af3f2ae55d9915505917191d
1 file changed
+175
benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py
new
+175
@@ -0,0 +1,175 @@
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 current_platform
11
+from vllm.utils.argparse_utils import FlexibleArgumentParser
12
+
13
+# Check if CPU MoE operations are available
14
+try:
15
+ from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
16
+except (ImportError, AttributeError) as e:
17
+ print("ERROR: CPU fused MoE operations are not available on this platform.")
18
+ print("This benchmark requires x86 CPU with proper vLLM CPU extensions compiled.")
19
+ print(
20
+ "The cpu_fused_moe kernel is typically available on Linux x86_64 "
21
+ "with AVX2/AVX512."
22
+ )
23
+ print(f"Import error: {e}")
24
+ sys.exit(1)
25
+
26
+# ISA selection following test_cpu_fused_moe.py pattern
27
+ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
28
+
29
+
30
+@torch.inference_mode()
31
+def main(
32
+ batch_size: int,
33
+ expert_num: int,
34
+ hidden_size: int,
35
+ intermediate_size: int,
36
+ topk_num: int,
37
+ use_bias: bool = False,
38
+ dtype: torch.dtype = torch.bfloat16,
39
+ activation: str = "silu",
40
+ isa: str = "vec",
41
+ seed: int = 0,
42
+ iters: int = 20,
43
+) -> None:
44
+ current_platform.seed_everything(seed)
45
+ # up_dim = 2 * intermediate_size for gate + up projection
46
+ up_dim = 2 * intermediate_size
47
+
48
+ input_tensor = torch.randn((batch_size, hidden_size), dtype=dtype) / (
49
+ 0.5 * hidden_size**0.5
50
+ )
51
+
52
+ w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
53
+ 0.5 * hidden_size**0.5
54
+ )
55
+ w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
56
+ 0.5 * intermediate_size**0.5
57
+ )
58
+
59
+ w13_bias = None
60
+ w2_bias = None
61
+ if use_bias:
62
+ w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
63
+ w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
64
+ 0.5 * hidden_size**0.5
65
+ )
66
+
67
+ router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
68
+ score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
69
+ topk_weights, topk_ids = torch.topk(score, topk_num)
70
+ topk_ids = topk_ids.to(torch.int32)
71
+
72
+ packed_w13 = cpu_prepack_moe_weight(w13, isa)
73
+ packed_w2 = cpu_prepack_moe_weight(w2, isa)
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
+ input_tensor,
81
+ packed_w13,
82
+ packed_w2,
83
+ w13_bias,
84
+ w2_bias,
85
+ topk_weights,
86
+ topk_ids,
87
+ activation,
88
+ isa,
89
+ )
90
+ end_time = time.perf_counter_ns()
91
+ times.append((end_time - start_time) / 1e6)
92
+ return times
93
+
94
+ # warmup
95
+ run_benchmark(5)
96
+ # benchmark
97
+ times = run_benchmark(iters)
98
+
99
+ if not times:
100
+ print("No iterations to measure. Set --iters > 0.")
101
+ return
102
+
103
+ time_min = min(times)
104
+ time_max = max(times)
105
+ time_mean = np.mean(times)
106
+ time_std = np.std(times)
107
+
108
+ print("\tmin (ms) = ", time_min)
109
+ print("\tmax (ms) = ", time_max)
110
+ print("\tmean (ms) = ", time_mean)
111
+ print("\tstd = ", time_std)
112
+ print("\tmedian (ms) = ", np.median(times))
113
+
114
+ # Calculate throughput metrics
115
+ # FLOPs estimation: 2 * batch * topk * (hidden * up_dim + intermediate * hidden)
116
+ flops_per_token = (
117
+ 2 * topk_num * (hidden_size * up_dim + intermediate_size * hidden_size)
118
+ )
119
+ total_flops = batch_size * flops_per_token
120
+ tflops = total_flops / (time_mean * 1e-3) / 1e12
121
+ print(f"\tthroughput (TFLOP/s) = {tflops:.4f}")
122
+
123
+
124
+if __name__ == "__main__":
125
+ parser = FlexibleArgumentParser(description="Benchmark the CPU fused MoE kernel.")
126
+ parser.add_argument("--batch-size", type=int, default=64)
127
+ parser.add_argument("--expert-num", type=int, default=8)
128
+ parser.add_argument("--hidden-size", type=int, default=2880)
129
+ parser.add_argument("--intermediate-size", type=int, default=2880)
130
+ parser.add_argument(
131
+ "--topk-num",
132
+ type=int,
133
+ default=None,
134
+ help="Number of experts to route each token to (default: expert_num // 2)",
135
+ )
136
+ parser.add_argument("--use-bias", action="store_true")
137
+ parser.add_argument(
138
+ "--activation",
139
+ type=str,
140
+ choices=["silu", "swigluoai"],
141
+ default="silu",
142
+ help="Activation function",
143
+ )
144
+ parser.add_argument(
145
+ "--isa",
146
+ type=str,
147
+ choices=ISA_CHOICES,
148
+ default=ISA_CHOICES[0],
149
+ help=f"ISA to use (available: {ISA_CHOICES})",
150
+ )
151
+ parser.add_argument("--seed", type=int, default=0)
152
+ parser.add_argument("--iters", type=int, default=20)
153
+
154
+ args = parser.parse_args()
155
+
156
+ # Default topk_num to expert_num // 2, minimum 1
157
+ topk_num = (
158
+ args.topk_num if args.topk_num is not None else max(args.expert_num // 2, 1)
159
+ )
160
+
161
+ print(args)
162
+
163
+ main(
164
+ batch_size=args.batch_size,
165
+ expert_num=args.expert_num,
166
+ hidden_size=args.hidden_size,
167
+ intermediate_size=args.intermediate_size,
168
+ topk_num=topk_num,
169
+ use_bias=args.use_bias,
170
+ dtype=torch.bfloat16, # Following test_cpu_fused_moe.py
171
+ activation=args.activation,
172
+ isa=args.isa,
173
+ seed=args.seed,
174
+ iters=args.iters,
175
+ )