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 14, 05:30 UTC
Add CUTLASS FP8 MOE benchmark scripts and kernel config (#25302)
Signed-off-by: Chenxi Yang <cxyang@fb.com> Co-authored-by: Chenxi Yang <cxyang@fb.com>
Chenxi Yang committed
Sep 23, 2025 at 17:07 UTC
0d235b874aae47b3214b3354bb70fcefb91c4fce
2 files changed
+529
benchmarks/kernels/benchmark_cutlass_moe_fp8.py
new
+406
@@ -0,0 +1,406 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
5
+kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
6
+but use different quantization strategies and backends.
7
+"""
8
+
9
+import nvtx
10
+import torch
11
+
12
+from vllm import _custom_ops as ops
13
+from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
14
+from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
15
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
16
+from vllm.platforms import current_platform
17
+from vllm.utils import FlexibleArgumentParser
18
+
19
+# Weight shapes for different models: [num_experts, topk, hidden_size,
20
+# intermediate_size]
21
+WEIGHT_SHAPES_MOE = {
22
+ "mixtral-8x7b": [
23
+ [8, 2, 4096, 14336],
24
+ ],
25
+ "deepseek-v2": [
26
+ [160, 6, 5120, 12288],
27
+ ],
28
+ "custom-small": [
29
+ [8, 2, 2048, 7168],
30
+ ],
31
+ "glm45-fp8": [
32
+ [128, 8, 4096, 1408],
33
+ ],
34
+ "Llama-4-Maverick-17B-128E-Instruct-FP8": [
35
+ [128, 1, 5120, 8192],
36
+ ],
37
+}
38
+
39
+DEFAULT_MODELS = [
40
+ "mixtral-8x7b",
41
+]
42
+
43
+DEFAULT_BATCH_SIZES = [4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]
44
+DEFAULT_TP_SIZES = [1]
45
+
46
+PER_ACT_TOKEN_OPTS = [False, True]
47
+PER_OUT_CH_OPTS = [False, True]
48
+
49
+FP8_DTYPE = current_platform.fp8_dtype()
50
+
51
+
52
+def bench_run(
53
+ results: list,
54
+ model: str,
55
+ num_experts: int,
56
+ topk: int,
57
+ per_act_token: bool,
58
+ per_out_ch: bool,
59
+ mkn: tuple[int, int, int],
60
+):
61
+ (m, k, n) = mkn
62
+
63
+ dtype = torch.half
64
+ device = "cuda"
65
+
66
+ # Create input activations
67
+ a = torch.randn((m, k), device=device, dtype=dtype) / 10
68
+
69
+ # Create weights
70
+ w1 = torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 10
71
+ w2 = torch.randn((num_experts, k, n), device=device, dtype=dtype) / 10
72
+
73
+ # Create FP8 quantized weights and scales for both kernels
74
+ w1_fp8q = torch.empty((num_experts, 2 * n, k), device=device, dtype=FP8_DTYPE)
75
+ w2_fp8q = torch.empty((num_experts, k, n), device=device, dtype=FP8_DTYPE)
76
+
77
+ # Create scales based on quantization strategy
78
+ if per_out_ch:
79
+ # Per-channel quantization
80
+ w1_scale = torch.empty(
81
+ (num_experts, 2 * n, 1), device=device, dtype=torch.float32
82
+ )
83
+ w2_scale = torch.empty((num_experts, k, 1), device=device, dtype=torch.float32)
84
+ else:
85
+ # Per-tensor quantization
86
+ w1_scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)
87
+ w2_scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)
88
+
89
+ # Quantize weights
90
+ for expert in range(num_experts):
91
+ if per_out_ch:
92
+ # Per-channel quantization - not yet implemented properly
93
+ # For now, fall back to per-tensor quantization
94
+ w1_fp8q[expert], w1_scale_temp = ops.scaled_fp8_quant(w1[expert])
95
+ w2_fp8q[expert], w2_scale_temp = ops.scaled_fp8_quant(w2[expert])
96
+ # Expand scalar scales to the expected per-channel shape
97
+ w1_scale[expert] = w1_scale_temp.expand(2 * n, 1)
98
+ w2_scale[expert] = w2_scale_temp.expand(k, 1)
99
+ else:
100
+ # Per-tensor quantization
101
+ w1_fp8q[expert], w1_scale_temp = ops.scaled_fp8_quant(w1[expert])
102
+ w2_fp8q[expert], w2_scale_temp = ops.scaled_fp8_quant(w2[expert])
103
+ # Store scalar scales in [1, 1] tensors
104
+ w1_scale[expert, 0, 0] = w1_scale_temp
105
+ w2_scale[expert, 0, 0] = w2_scale_temp
106
+
107
+ # Prepare weights for CUTLASS (no transpose needed)
108
+ w1_fp8q_cutlass = w1_fp8q # Keep original [E, 2N, K]
109
+ w2_fp8q_cutlass = w2_fp8q # Keep original [E, K, N]
110
+
111
+ # Create router scores and get topk
112
+ score = torch.randn((m, num_experts), device=device, dtype=dtype)
113
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
114
+
115
+ # WORKAROUND: CUTLASS MoE FP8 has issues with per-token quantization
116
+ # Force per-tensor quantization for all cases to match working e2e setup
117
+ a1_scale = torch.full((), 1e-2, device=device, dtype=torch.float32)
118
+ a2_scale = torch.full((), 1e-2, device=device, dtype=torch.float32)
119
+
120
+ # Force per-tensor quantization for all cases
121
+ per_act_token = False
122
+
123
+ # Create stride tensors for CUTLASS
124
+ ab_strides1 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
125
+ ab_strides2 = torch.full((num_experts,), n, dtype=torch.int64, device=device)
126
+ c_strides1 = torch.full((num_experts,), 2 * n, dtype=torch.int64, device=device)
127
+ c_strides2 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
128
+
129
+ def run_triton_moe(
130
+ a: torch.Tensor,
131
+ w1: torch.Tensor,
132
+ w2: torch.Tensor,
133
+ topk_weights: torch.Tensor,
134
+ topk_ids: torch.Tensor,
135
+ w1_scale: torch.Tensor,
136
+ w2_scale: torch.Tensor,
137
+ a1_scale: torch.Tensor,
138
+ a2_scale: torch.Tensor,
139
+ num_repeats: int,
140
+ ):
141
+ quant_config = fp8_w8a8_moe_quant_config(
142
+ w1_scale=w1_scale,
143
+ w2_scale=w2_scale,
144
+ a1_scale=a1_scale,
145
+ a2_scale=a2_scale,
146
+ per_act_token_quant=per_act_token,
147
+ per_out_ch_quant=per_out_ch,
148
+ )
149
+
150
+ for _ in range(num_repeats):
151
+ fused_experts(
152
+ a,
153
+ w1,
154
+ w2,
155
+ topk_weights,
156
+ topk_ids,
157
+ quant_config=quant_config,
158
+ )
159
+
160
+ def run_cutlass_moe_fp8(
161
+ a: torch.Tensor,
162
+ w1: torch.Tensor,
163
+ w2: torch.Tensor,
164
+ topk_weights: torch.Tensor,
165
+ topk_ids: torch.Tensor,
166
+ ab_strides1: torch.Tensor,
167
+ ab_strides2: torch.Tensor,
168
+ c_strides1: torch.Tensor,
169
+ c_strides2: torch.Tensor,
170
+ w1_scale: torch.Tensor,
171
+ w2_scale: torch.Tensor,
172
+ a1_scale: torch.Tensor,
173
+ a2_scale: torch.Tensor,
174
+ num_repeats: int,
175
+ ):
176
+ quant_config = fp8_w8a8_moe_quant_config(
177
+ w1_scale=w1_scale,
178
+ w2_scale=w2_scale,
179
+ a1_scale=a1_scale,
180
+ a2_scale=a2_scale,
181
+ per_act_token_quant=per_act_token,
182
+ per_out_ch_quant=per_out_ch,
183
+ )
184
+
185
+ for _ in range(num_repeats):
186
+ with nvtx.annotate("cutlass_moe_fp8", color="blue"):
187
+ cutlass_moe_fp8(
188
+ a=a,
189
+ w1_q=w1,
190
+ w2_q=w2,
191
+ topk_weights=topk_weights,
192
+ topk_ids=topk_ids,
193
+ ab_strides1=ab_strides1,
194
+ ab_strides2=ab_strides2,
195
+ c_strides1=c_strides1,
196
+ c_strides2=c_strides2,
197
+ quant_config=quant_config,
198
+ activation="silu",
199
+ global_num_experts=num_experts,
200
+ )
201
+
202
+ # Pre-create quantization config to avoid creating it inside CUDA graph
203
+ quant_config = fp8_w8a8_moe_quant_config(
204
+ w1_scale=w1_scale,
205
+ w2_scale=w2_scale,
206
+ a1_scale=a1_scale,
207
+ a2_scale=a2_scale,
208
+ per_act_token_quant=per_act_token,
209
+ per_out_ch_quant=per_out_ch,
210
+ )
211
+
212
+ # Create CUDA graphs for CUTLASS (match benchmark_moe.py pattern exactly)
213
+ cutlass_stream = torch.cuda.Stream()
214
+ cutlass_graph = torch.cuda.CUDAGraph()
215
+ with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
216
+ # Capture 10 invocations like benchmark_moe.py
217
+ for _ in range(10):
218
+ cutlass_moe_fp8(
219
+ a=a,
220
+ w1_q=w1_fp8q_cutlass,
221
+ w2_q=w2_fp8q_cutlass,
222
+ topk_weights=topk_weights,
223
+ topk_ids=topk_ids,
224
+ ab_strides1=ab_strides1,
225
+ ab_strides2=ab_strides2,
226
+ c_strides1=c_strides1,
227
+ c_strides2=c_strides2,
228
+ quant_config=quant_config,
229
+ activation="silu",
230
+ global_num_experts=num_experts,
231
+ )
232
+ torch.cuda.synchronize()
233
+
234
+ # Create CUDA graphs for Triton (match benchmark_moe.py pattern exactly)
235
+ triton_stream = torch.cuda.Stream()
236
+ triton_graph = torch.cuda.CUDAGraph()
237
+ with torch.cuda.graph(triton_graph, stream=triton_stream):
238
+ # Capture 10 invocations like benchmark_moe.py
239
+ for _ in range(10):
240
+ fused_experts(
241
+ a,
242
+ w1_fp8q,
243
+ w2_fp8q,
244
+ topk_weights,
245
+ topk_ids,
246
+ quant_config=quant_config,
247
+ )
248
+ torch.cuda.synchronize()
249
+
250
+ def bench_cuda_graph(graph, num_warmup=5, num_iters=100):
251
+ """Benchmark CUDA graph using events like benchmark_moe.py"""
252
+ # Warmup
253
+ for _ in range(num_warmup):
254
+ graph.replay()
255
+ torch.cuda.synchronize()
256
+
257
+ # Timing
258
+ start_event = torch.cuda.Event(enable_timing=True)
259
+ end_event = torch.cuda.Event(enable_timing=True)
260
+
261
+ latencies = []
262
+ for _ in range(num_iters):
263
+ torch.cuda.synchronize()
264
+ start_event.record()
265
+ graph.replay()
266
+ end_event.record()
267
+ end_event.synchronize()
268
+ latencies.append(start_event.elapsed_time(end_event))
269
+
270
+ # Divide by 10 since graph contains 10 calls
271
+ return sum(latencies) / (num_iters * 10)
272
+
273
+ # Benchmark parameters
274
+ num_warmup = 5
275
+ num_iters = 100
276
+
277
+ # Benchmark only CUDA graphs (more reliable and faster)
278
+ # Benchmark Triton MoE with CUDA graphs
279
+ triton_graph_time = bench_cuda_graph(
280
+ triton_graph, num_warmup=num_warmup, num_iters=num_iters
281
+ )
282
+
283
+ # Benchmark CUTLASS MoE with CUDA graphs
284
+ cutlass_graph_time = bench_cuda_graph(
285
+ cutlass_graph, num_warmup=num_warmup, num_iters=num_iters
286
+ )
287
+
288
+ # Convert ms to us and return results
289
+ triton_time_us = triton_graph_time * 1000
290
+ cutlass_time_us = cutlass_graph_time * 1000
291
+
292
+ return {
293
+ "batch_size": m,
294
+ "triton_time_us": triton_time_us,
295
+ "cutlass_time_us": cutlass_time_us,
296
+ }
297
+
298
+
299
+def main(args):
300
+ print("Benchmarking models:")
301
+ for i, model in enumerate(args.models):
302
+ print(f"[{i}] {model}")
303
+
304
+ all_results = []
305
+
306
+ for model in args.models:
307
+ for tp in args.tp_sizes:
308
+ for layer in WEIGHT_SHAPES_MOE[model]:
309
+ num_experts = layer[0]
310
+ topk = layer[1]
311
+ size_k = layer[2]
312
+ size_n = layer[3] // tp
313
+
314
+ if len(args.limit_k) > 0 and size_k not in args.limit_k:
315
+ continue
316
+
317
+ if len(args.limit_n) > 0 and size_n not in args.limit_n:
318
+ continue
319
+
320
+ for per_act_token in args.per_act_token_opts:
321
+ for per_out_ch in args.per_out_ch_opts:
322
+ print(
323
+ f"\n=== {model}, experts={num_experts}, topk={topk},"
324
+ f"per_act={per_act_token}, per_out_ch={per_out_ch} ==="
325
+ )
326
+
327
+ config_results = []
328
+ for size_m in args.batch_sizes:
329
+ mkn = (size_m, size_k, size_n)
330
+ result = bench_run(
331
+ [], # Not used anymore
332
+ model,
333
+ num_experts,
334
+ topk,
335
+ per_act_token,
336
+ per_out_ch,
337
+ mkn,
338
+ )
339
+ if result:
340
+ config_results.append(result)
341
+
342
+ # Print results table for this configuration
343
+ if config_results:
344
+ print(
345
+ f"\n{'Batch Size':<12}"
346
+ f"{'Triton (us)':<15}"
347
+ f"{'CUTLASS (us)':<15}"
348
+ )
349
+ print("-" * 45)
350
+ for result in config_results:
351
+ print(
352
+ f"{result['batch_size']:<12}"
353
+ f"{result['triton_time_us']:<15.2f}"
354
+ f"{result['cutlass_time_us']:<15.2f}"
355
+ )
356
+
357
+ all_results.extend(config_results)
358
+
359
+ print(f"\nTotal benchmarks completed: {len(all_results)}")
360
+
361
+
362
+if __name__ == "__main__":
363
+ parser = FlexibleArgumentParser(
364
+ description="""Benchmark CUTLASS FP8 MOE vs Triton FP8 FUSED MOE
365
+ across specified models/shapes/batches
366
+
367
+ Example usage:
368
+ python benchmark_cutlass_moe_fp8.py \
369
+ --model "Llama-4-Maverick-17B-128E-Instruct-FP8" \
370
+ --tp-sizes 8 \
371
+ --batch-size 2 4 8 \
372
+ --per-act-token-opts false \
373
+ --per-out-ch-opts false
374
+
375
+ """
376
+ )
377
+ parser.add_argument(
378
+ "--models",
379
+ nargs="+",
380
+ type=str,
381
+ default=DEFAULT_MODELS,
382
+ choices=WEIGHT_SHAPES_MOE.keys(),
383
+ )
384
+ parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES)
385
+ parser.add_argument(
386
+ "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
387
+ )
388
+ parser.add_argument("--limit-k", nargs="+", type=int, default=[])
389
+ parser.add_argument("--limit-n", nargs="+", type=int, default=[])
390
+ parser.add_argument(
391
+ "--per-act-token-opts",
392
+ nargs="+",
393
+ type=lambda x: x.lower() == "true",
394
+ default=[False, True],
395
+ help="Per-activation token quantization options (true/false)",
396
+ )
397
+ parser.add_argument(
398
+ "--per-out-ch-opts",
399
+ nargs="+",
400
+ type=lambda x: x.lower() == "true",
401
+ default=[False, True],
402
+ help="Per-output channel quantization options (true/false)",
403
+ )
404
+
405
+ args = parser.parse_args()
406
+ main(args)
vllm/model_executor/layers/fused_moe/configs/E=128,N=1024,device_name=NVIDIA_H100,dtype=fp8_w8a8.json
new
+123
@@ -0,0 +1,123 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "2": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 256,
7
+ "GROUP_SIZE_M": 16,
8
+ "num_warps": 4,
9
+ "num_stages": 4
10
+ },
11
+ "4": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 32,
14
+ "BLOCK_SIZE_K": 256,
15
+ "GROUP_SIZE_M": 64,
16
+ "num_warps": 4,
17
+ "num_stages": 4
18
+ },
19
+ "8": {
20
+ "BLOCK_SIZE_M": 64,
21
+ "BLOCK_SIZE_N": 64,
22
+ "BLOCK_SIZE_K": 256,
23
+ "GROUP_SIZE_M": 1,
24
+ "num_warps": 8,
25
+ "num_stages": 3
26
+ },
27
+ "16": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 32,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 32,
32
+ "num_warps": 4,
33
+ "num_stages": 3
34
+ },
35
+ "32": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 32,
38
+ "BLOCK_SIZE_K": 256,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 5
42
+ },
43
+ "64": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 256,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 8,
49
+ "num_stages": 3
50
+ },
51
+ "128": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 32,
54
+ "BLOCK_SIZE_K": 256,
55
+ "GROUP_SIZE_M": 16,
56
+ "num_warps": 4,
57
+ "num_stages": 5
58
+ },
59
+ "256": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 32,
62
+ "BLOCK_SIZE_K": 256,
63
+ "GROUP_SIZE_M": 16,
64
+ "num_warps": 4,
65
+ "num_stages": 5
66
+ },
67
+ "512": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 32,
70
+ "BLOCK_SIZE_K": 256,
71
+ "GROUP_SIZE_M": 16,
72
+ "num_warps": 4,
73
+ "num_stages": 5
74
+ },
75
+ "1024": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 32,
78
+ "BLOCK_SIZE_K": 256,
79
+ "GROUP_SIZE_M": 16,
80
+ "num_warps": 4,
81
+ "num_stages": 5
82
+ },
83
+ "2048": {
84
+ "BLOCK_SIZE_M": 32,
85
+ "BLOCK_SIZE_N": 128,
86
+ "BLOCK_SIZE_K": 256,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 8,
89
+ "num_stages": 3
90
+ },
91
+ "3072": {
92
+ "BLOCK_SIZE_M": 32,
93
+ "BLOCK_SIZE_N": 128,
94
+ "BLOCK_SIZE_K": 256,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 8,
97
+ "num_stages": 3
98
+ },
99
+ "4096": {
100
+ "BLOCK_SIZE_M": 64,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 4,
105
+ "num_stages": 3
106
+ },
107
+ "8192": {
108
+ "BLOCK_SIZE_M": 128,
109
+ "BLOCK_SIZE_N": 256,
110
+ "BLOCK_SIZE_K": 128,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 8,
113
+ "num_stages": 4
114
+ },
115
+ "16384": {
116
+ "BLOCK_SIZE_M": 128,
117
+ "BLOCK_SIZE_N": 256,
118
+ "BLOCK_SIZE_K": 128,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 8,
121
+ "num_stages": 4
122
+ }
123
+}