| 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.benchmarks.lib.utils import default_vllm_config |
| 11 | from vllm.model_executor.custom_op import op_registry |
| 12 | from vllm.triton_utils import triton |
| 13 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 14 | from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed |
| 15 | |
| 16 | batch_size_range = [1, 16, 128] |
| 17 | seq_len_range = [1, 16, 64, 1024, 4096] |
| 18 | intermediate_size = [3072, 9728, 12288] |
| 19 | configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size)) |
| 20 | |
| 21 | |
| 22 | @default_vllm_config() |
| 23 | def benchmark_activation( |
| 24 | batch_size: int, |
| 25 | seq_len: int, |
| 26 | intermediate_size: int, |
| 27 | provider: str, |
| 28 | func_name: str, |
| 29 | dtype: torch.dtype, |
| 30 | ): |
| 31 | device = "cuda" |
| 32 | num_tokens = batch_size * seq_len |
| 33 | dim = intermediate_size |
| 34 | set_random_seed(42) |
| 35 | torch.set_default_device(device) |
| 36 | |
| 37 | if func_name == "gelu_and_mul": |
| 38 | layer = op_registry[func_name](approximate="none") |
| 39 | elif func_name == "gelu_and_mul_tanh": |
| 40 | layer = op_registry["gelu_and_mul"](approximate="tanh") |
| 41 | elif func_name == "fatrelu_and_mul": |
| 42 | threshold = 0.5 |
| 43 | layer = op_registry[func_name](threshold) |
| 44 | else: |
| 45 | layer = op_registry[func_name]() |
| 46 | |
| 47 | x = torch.randn(num_tokens, dim, dtype=dtype, device=device) |
| 48 | compiled_layer = torch.compile(layer.forward_native) |
| 49 | |
| 50 | if provider == "custom": |
| 51 | fn = lambda: layer(x) |
| 52 | elif provider == "compiled": |
| 53 | fn = lambda: compiled_layer(x) |
| 54 | |
| 55 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 56 | fn, quantiles=[0.5, 0.2, 0.8] |
| 57 | ) |
| 58 | return ms, max_ms, min_ms |
| 59 | |
| 60 | |
| 61 | if __name__ == "__main__": |
| 62 | parser = FlexibleArgumentParser(description="Benchmark the custom activation op.") |
| 63 | parser.add_argument( |
| 64 | "--func-name", |
| 65 | type=str, |
| 66 | choices=[ |
| 67 | "mul_and_silu", |
| 68 | "silu_and_mul", |
| 69 | "gelu_and_mul", |
| 70 | "gelu_and_mul_tanh", |
| 71 | "fatrelu_and_mul", |
| 72 | "swigluoai_and_mul", |
| 73 | "gelu_new", |
| 74 | "gelu_fast", |
| 75 | "quick_gelu", |
| 76 | ], |
| 77 | default="silu_and_mul", |
| 78 | ) |
| 79 | parser.add_argument( |
| 80 | "--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16" |
| 81 | ) |
| 82 | args = parser.parse_args() |
| 83 | assert args |
| 84 | |
| 85 | func_name = args.func_name |
| 86 | dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype] |
| 87 | |
| 88 | perf_report = triton.testing.perf_report( |
| 89 | triton.testing.Benchmark( |
| 90 | x_names=["batch_size", "seq_len", "intermediate_size"], |
| 91 | x_vals=configs, |
| 92 | line_arg="provider", |
| 93 | line_vals=["custom", "compiled"], |
| 94 | line_names=["Custom OP", "Compiled"], |
| 95 | styles=[("blue", "-"), ("green", "-")], |
| 96 | ylabel="ms", |
| 97 | plot_name=f"{func_name}-op-performance", |
| 98 | args={}, |
| 99 | ) |
| 100 | ) |
| 101 | |
| 102 | perf_report( |
| 103 | lambda batch_size, seq_len, intermediate_size, provider: benchmark_activation( |
| 104 | batch_size, seq_len, intermediate_size, provider, func_name, dtype |
| 105 | ) |
| 106 | ).run(print_data=True) |