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
[Benchmark] add benchmark for custom activation op (#23908)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Jiangyun Zhu committed
Sep 7, 2025 at 11:12 UTC
77aec83b8c6e6c15a9b5c333a531c29eff0b61fc
1 file changed
+104
benchmarks/kernels/benchmark_activation.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# benchmark custom activation op performance
5
+import itertools
6
+
7
+import torch
8
+
9
+import vllm.model_executor.layers.activation # noqa F401
10
+from vllm.model_executor.custom_op import CustomOp
11
+from vllm.platforms import current_platform
12
+from vllm.triton_utils import triton
13
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
14
+
15
+batch_size_range = [1, 16, 32, 64, 128]
16
+seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
17
+intermediate_size = [3072, 9728, 12288]
18
+configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
19
+
20
+
21
+def benchmark_activation(
22
+ batch_size: int,
23
+ seq_len: int,
24
+ intermediate_size: int,
25
+ provider: str,
26
+ func_name: str,
27
+ dtype: torch.dtype,
28
+):
29
+ device = "cuda"
30
+ num_tokens = batch_size * seq_len
31
+ dim = intermediate_size
32
+ current_platform.seed_everything(42)
33
+ torch.set_default_device(device)
34
+
35
+ if func_name == "gelu_and_mul":
36
+ layer = CustomOp.op_registry[func_name](approximate="none")
37
+ elif func_name == "gelu_and_mul_tanh":
38
+ layer = CustomOp.op_registry["gelu_and_mul"](approximate="tanh")
39
+ elif func_name == "fatrelu_and_mul":
40
+ threshold = 0.5
41
+ layer = CustomOp.op_registry[func_name](threshold)
42
+ else:
43
+ layer = CustomOp.op_registry[func_name]()
44
+
45
+ x = torch.randn(num_tokens, dim, dtype=dtype, device=device)
46
+ compiled_layer = torch.compile(layer.forward_native)
47
+
48
+ if provider == "custom":
49
+ fn = lambda: layer(x)
50
+ elif provider == "compiled":
51
+ fn = lambda: compiled_layer(x)
52
+
53
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
54
+ fn, quantiles=[0.5, 0.2, 0.8]
55
+ )
56
+ return ms, max_ms, min_ms
57
+
58
+
59
+if __name__ == "__main__":
60
+ parser = FlexibleArgumentParser(description="Benchmark the custom activation op.")
61
+ parser.add_argument(
62
+ "--func-name",
63
+ type=str,
64
+ choices=[
65
+ "mul_and_silu",
66
+ "silu_and_mul",
67
+ "gelu_and_mul",
68
+ "gelu_and_mul_tanh",
69
+ "fatrelu_and_mul",
70
+ "swigluoai_and_mul",
71
+ "gelu_new",
72
+ "gelu_fast",
73
+ "quick_gelu",
74
+ ],
75
+ default="silu_and_mul",
76
+ )
77
+ parser.add_argument(
78
+ "--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16"
79
+ )
80
+ args = parser.parse_args()
81
+ assert args
82
+
83
+ func_name = args.func_name
84
+ dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
85
+
86
+ perf_report = triton.testing.perf_report(
87
+ triton.testing.Benchmark(
88
+ x_names=["batch_size", "seq_len", "intermediate_size"],
89
+ x_vals=configs,
90
+ line_arg="provider",
91
+ line_vals=["custom", "compiled"],
92
+ line_names=["Custom OP", "Compiled"],
93
+ styles=[("blue", "-"), ("green", "-")],
94
+ ylabel="ms",
95
+ plot_name=f"{func_name}-op-performance",
96
+ args={},
97
+ )
98
+ )
99
+
100
+ perf_report(
101
+ lambda batch_size, seq_len, intermediate_size, provider: benchmark_activation(
102
+ batch_size, seq_len, intermediate_size, provider, func_name, dtype
103
+ )
104
+ ).run(print_data=True)