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