| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both |
| 5 | # eager and under torch.compile (Inductor fuses relu+square into one kernel). |
| 6 | |
| 7 | import itertools |
| 8 | |
| 9 | import torch |
| 10 | import torch.nn.functional as F |
| 11 | |
| 12 | import vllm.model_executor.layers.activation # noqa: F401 |
| 13 | from vllm.benchmarks.lib.utils import default_vllm_config |
| 14 | from vllm.triton_utils import triton |
| 15 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 16 | from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed |
| 17 | |
| 18 | # Capped so the largest tensor stays under 2**31 elements: the shared activation |
| 19 | # kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which |
| 20 | # overflows for tensors with >2**32 elements. Realistic token counts are well |
| 21 | # below this; the kernel-vs-native gap is already clear at these sizes. |
| 22 | batch_size_range = [1, 16, 128] |
| 23 | seq_len_range = [1, 16, 64, 1024] |
| 24 | intermediate_size = [3072, 9728, 12288] |
| 25 | configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size)) |
| 26 | |
| 27 | |
| 28 | @default_vllm_config() |
| 29 | def benchmark_relu_squared( |
| 30 | batch_size: int, |
| 31 | seq_len: int, |
| 32 | intermediate_size: int, |
| 33 | provider: str, |
| 34 | dtype: torch.dtype, |
| 35 | ): |
| 36 | device = "cuda" |
| 37 | num_tokens = batch_size * seq_len |
| 38 | set_random_seed(42) |
| 39 | torch.set_default_device(device) |
| 40 | |
| 41 | x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device) |
| 42 | out = torch.empty_like(x) |
| 43 | |
| 44 | def native(x: torch.Tensor) -> torch.Tensor: |
| 45 | return torch.square(F.relu(x)) |
| 46 | |
| 47 | # Verify the custom kernel matches the native implementation before timing. |
| 48 | ref = native(x) |
| 49 | torch.ops._C.relu_squared(out, x) |
| 50 | torch.testing.assert_close(out, ref) |
| 51 | |
| 52 | if provider == "custom": |
| 53 | # Custom CUDA kernel — single fused kernel. |
| 54 | fn = lambda: torch.ops._C.relu_squared(out, x) |
| 55 | elif provider == "native": |
| 56 | # forward_native, eager — relu and square as separate ops. |
| 57 | fn = lambda: native(x) |
| 58 | elif provider == "native_compiled": |
| 59 | # forward_native under torch.compile — Inductor fuses relu+square. |
| 60 | # This is the real production baseline (custom ops are off when |
| 61 | # Inductor is enabled), so it is the comparison reviewers care about. |
| 62 | compiled = torch.compile(native) |
| 63 | compiled(x) # warm up / trigger compilation before timing |
| 64 | fn = lambda: compiled(x) |
| 65 | |
| 66 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 67 | fn, quantiles=[0.5, 0.2, 0.8] |
| 68 | ) |
| 69 | return ms, max_ms, min_ms |
| 70 | |
| 71 | |
| 72 | if __name__ == "__main__": |
| 73 | parser = FlexibleArgumentParser( |
| 74 | description="Benchmark ReLUSquaredActivation: custom kernel vs native." |
| 75 | ) |
| 76 | parser.add_argument( |
| 77 | "--dtype", |
| 78 | type=str, |
| 79 | choices=["half", "bfloat16", "float"], |
| 80 | default="bfloat16", |
| 81 | ) |
| 82 | args = parser.parse_args() |
| 83 | |
| 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", "native_compiled", "native"], |
| 92 | line_names=[ |
| 93 | "Custom Kernel", |
| 94 | "Native (torch.compile)", |
| 95 | "Native (eager)", |
| 96 | ], |
| 97 | styles=[("blue", "-"), ("green", "-"), ("red", "-")], |
| 98 | ylabel="ms", |
| 99 | plot_name="relu_squared-eager-performance", |
| 100 | args={}, |
| 101 | ) |
| 102 | ) |
| 103 | |
| 104 | perf_report( |
| 105 | lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared( |
| 106 | batch_size, seq_len, intermediate_size, provider, dtype |
| 107 | ) |
| 108 | ).run(print_data=True) |