| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # This script benchmarks the mrope kernel (mainly for Qwen2VL and Qwen2.5VL models). |
| 5 | # It generates test data, runs benchmarks, and saves results to a CSV file. |
| 6 | # |
| 7 | # The CSV file (named with current date/time) contains these columns: |
| 8 | # model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position, |
| 9 | # is_neox_style, rope_parameters, dtype, torch_mean, torch_median, torch_p99, |
| 10 | # torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max, |
| 11 | # speedup |
| 12 | # |
| 13 | # == Usage Examples == |
| 14 | # |
| 15 | # Single model benchmark: |
| 16 | # python3 benchmark_mrope.py --model-name Qwen/Qwen2-VL-7B-Instruct --tp-size 1 \ |
| 17 | # --warmup-iter 10 --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024 |
| 18 | # |
| 19 | # All models benchmark: |
| 20 | # python3 benchmark_mrope.py --model-name "" --tp-size 1 --warmup-iter 10 \ |
| 21 | # --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024 |
| 22 | # |
| 23 | # All models with different TP sizes: |
| 24 | # python3 benchmark_mrope.py --model-name "" --tp-size 1 2 4 8 --warmup-iter 10 \ |
| 25 | # --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024 |
| 26 | # |
| 27 | # All models with different token counts: |
| 28 | # python3 benchmark_mrope.py --model-name "" --tp-size 1 --warmup-iter 10 \ |
| 29 | # --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024 4096 16384 |
| 30 | import csv |
| 31 | import os |
| 32 | import time |
| 33 | from datetime import datetime |
| 34 | from typing import Any |
| 35 | |
| 36 | import numpy as np |
| 37 | import torch |
| 38 | |
| 39 | from vllm.benchmarks.lib.utils import default_vllm_config |
| 40 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 41 | from vllm.transformers_utils.config import get_config |
| 42 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 43 | from vllm.utils.torch_utils import set_random_seed |
| 44 | |
| 45 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 46 | |
| 47 | |
| 48 | def generate_test_data( |
| 49 | num_tokens: int, |
| 50 | num_q_heads: int, |
| 51 | num_kv_heads: int, |
| 52 | head_size: int, |
| 53 | max_position_embeddings: int, |
| 54 | dtype: torch.dtype, |
| 55 | device: torch.device, |
| 56 | ): |
| 57 | """Generate test data for given configuration.""" |
| 58 | # Create 2D positions (3, num_tokens) for multimodal case |
| 59 | positions = torch.randint( |
| 60 | 0, max_position_embeddings // 4, (3, num_tokens), device=device |
| 61 | ) |
| 62 | |
| 63 | # Create query and key tensors |
| 64 | query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device) |
| 65 | key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device) |
| 66 | |
| 67 | return positions, query, key |
| 68 | |
| 69 | |
| 70 | def calculate_stats(times: list[float]) -> dict[str, float]: |
| 71 | """Calculate statistics from a list of times.""" |
| 72 | times_array = np.array(times) |
| 73 | return { |
| 74 | "mean": np.mean(times_array), |
| 75 | "median": np.median(times_array), |
| 76 | "p99": np.percentile(times_array, 99), |
| 77 | "min": np.min(times_array), |
| 78 | "max": np.max(times_array), |
| 79 | } |
| 80 | |
| 81 | |
| 82 | @default_vllm_config() |
| 83 | def benchmark_mrope( |
| 84 | model_name: str, |
| 85 | num_tokens: int, |
| 86 | head_dim: int, |
| 87 | tp_size: int, |
| 88 | num_heads: int, |
| 89 | num_kv_heads: int, |
| 90 | max_position: int = 8192, |
| 91 | is_neox_style: bool = True, |
| 92 | rope_parameters: dict[str, Any] | None = None, |
| 93 | dtype: torch.dtype = torch.bfloat16, |
| 94 | seed: int = 0, |
| 95 | warmup_iter: int = 10, |
| 96 | benchmark_iter: int = 100, |
| 97 | csv_writer=None, |
| 98 | ): |
| 99 | set_random_seed(seed) |
| 100 | torch.set_default_device(device) |
| 101 | # the parameters to compute the q k v size based on tp_size |
| 102 | mrope_helper_class = get_rope( |
| 103 | head_size=head_dim, |
| 104 | max_position=max_position, |
| 105 | is_neox_style=is_neox_style, |
| 106 | rope_parameters=rope_parameters, |
| 107 | dtype=dtype, |
| 108 | ).to(device=device) |
| 109 | |
| 110 | print(80 * "=") |
| 111 | print( |
| 112 | f"Evaluating model: {model_name} " |
| 113 | f"with tp_size: {tp_size} " |
| 114 | f"and num_tokens: {num_tokens}, " |
| 115 | f"dtype: {dtype}" |
| 116 | ) |
| 117 | |
| 118 | # create q k v input tensors |
| 119 | # create rotary pos emb input tensors |
| 120 | positions, query, key = generate_test_data( |
| 121 | num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device |
| 122 | ) |
| 123 | |
| 124 | # Warm up |
| 125 | for _ in range(warmup_iter): |
| 126 | mrope_helper_class.forward_native( |
| 127 | positions, |
| 128 | query.clone(), |
| 129 | key.clone(), |
| 130 | ) |
| 131 | |
| 132 | mrope_helper_class.forward_cuda( |
| 133 | positions, |
| 134 | query.clone(), |
| 135 | key.clone(), |
| 136 | ) |
| 137 | |
| 138 | torch.accelerator.synchronize() |
| 139 | |
| 140 | # Time reference implementation |
| 141 | torch_times = [] |
| 142 | for _ in range(benchmark_iter): |
| 143 | query_clone = query.clone() |
| 144 | key_clone = key.clone() |
| 145 | torch.accelerator.synchronize() |
| 146 | start_time = time.time() |
| 147 | |
| 148 | mrope_helper_class.forward_native( |
| 149 | positions, |
| 150 | query_clone, |
| 151 | key_clone, |
| 152 | ) |
| 153 | |
| 154 | torch.accelerator.synchronize() |
| 155 | torch_times.append(time.time() - start_time) |
| 156 | |
| 157 | # Time triton kernel implementation |
| 158 | triton_times = [] |
| 159 | for _ in range(benchmark_iter): |
| 160 | query_clone = query.clone() |
| 161 | key_clone = key.clone() |
| 162 | torch.accelerator.synchronize() |
| 163 | start_time = time.time() |
| 164 | mrope_helper_class.forward_cuda( |
| 165 | positions, |
| 166 | query_clone, |
| 167 | key_clone, |
| 168 | ) |
| 169 | torch.accelerator.synchronize() |
| 170 | triton_times.append(time.time() - start_time) |
| 171 | |
| 172 | # Calculate statistics |
| 173 | torch_stats = calculate_stats(torch_times) |
| 174 | triton_stats = calculate_stats(triton_times) |
| 175 | print(f"\nPerformance for config ({num_tokens}, {num_heads}, {num_kv_heads}):") |
| 176 | |
| 177 | print( |
| 178 | f"Torch implementation: " |
| 179 | f"mean={torch_stats['mean']:.8f}s, " |
| 180 | f"median={torch_stats['median']:.8f}s, " |
| 181 | f"p99={torch_stats['p99']:.8f}s" |
| 182 | ) |
| 183 | |
| 184 | print( |
| 185 | f"Triton implementation: " |
| 186 | f"mean={triton_stats['mean']:.8f}s, " |
| 187 | f"median={triton_stats['median']:.8f}s, " |
| 188 | f"p99={triton_stats['p99']:.8f}s" |
| 189 | ) |
| 190 | |
| 191 | print( |
| 192 | f"Triton Speedup over Torch: {torch_stats['mean'] / triton_stats['mean']:.8f}x" |
| 193 | ) |
| 194 | |
| 195 | # Write to CSV |
| 196 | if csv_writer: |
| 197 | row = [ |
| 198 | model_name, |
| 199 | tp_size, |
| 200 | num_tokens, |
| 201 | num_heads, |
| 202 | num_kv_heads, |
| 203 | head_dim, |
| 204 | max_position, |
| 205 | is_neox_style, |
| 206 | str(rope_parameters), |
| 207 | str(dtype).split(".")[-1], |
| 208 | torch_stats["mean"], |
| 209 | torch_stats["median"], |
| 210 | torch_stats["p99"], |
| 211 | torch_stats["min"], |
| 212 | torch_stats["max"], |
| 213 | triton_stats["mean"], |
| 214 | triton_stats["median"], |
| 215 | triton_stats["p99"], |
| 216 | triton_stats["min"], |
| 217 | triton_stats["max"], |
| 218 | torch_stats["mean"] / triton_stats["mean"], # speedup |
| 219 | ] |
| 220 | csv_writer.writerow(row) |
| 221 | |
| 222 | return torch_stats, triton_stats |
| 223 | |
| 224 | |
| 225 | if __name__ == "__main__": |
| 226 | parser = FlexibleArgumentParser( |
| 227 | description="Benchmark the rotary embedding kernels." |
| 228 | ) |
| 229 | parser.add_argument("--model-name", type=str, default="") |
| 230 | parser.add_argument("--tp-size", type=int, default=1) |
| 231 | parser.add_argument("--warmup-iter", type=int, default=10) |
| 232 | parser.add_argument("--benchmark-iter", type=int, default=100) |
| 233 | parser.add_argument("--dtype", type=str, choices=["bfloat16"], default="bfloat16") |
| 234 | parser.add_argument("--seed", type=int, default=0) |
| 235 | parser.add_argument("--num-tokens", type=int, nargs="+", required=False) |
| 236 | parser.add_argument("--trust-remote-code", action="store_true") |
| 237 | parser.add_argument("--output-csv", type=str, default="mrope_benchmark_results.csv") |
| 238 | args = parser.parse_args() |
| 239 | print(args) |
| 240 | |
| 241 | # Create CSV file for results |
| 242 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 243 | csv_filename = f"{os.path.splitext(args.output_csv)[0]}_{timestamp}.csv" |
| 244 | |
| 245 | with open(csv_filename, "w", newline="") as csvfile: |
| 246 | csv_writer = csv.writer(csvfile) |
| 247 | # Write header |
| 248 | header = [ |
| 249 | "model_name", |
| 250 | "tp_size", |
| 251 | "num_tokens", |
| 252 | "num_heads", |
| 253 | "num_kv_heads", |
| 254 | "head_dim", |
| 255 | "max_position", |
| 256 | "is_neox_style", |
| 257 | "rope_parameters", |
| 258 | "dtype", |
| 259 | "torch_mean", |
| 260 | "torch_median", |
| 261 | "torch_p99", |
| 262 | "torch_min", |
| 263 | "torch_max", |
| 264 | "triton_mean", |
| 265 | "triton_median", |
| 266 | "triton_p99", |
| 267 | "triton_min", |
| 268 | "triton_max", |
| 269 | "speedup", |
| 270 | ] |
| 271 | csv_writer.writerow(header) |
| 272 | |
| 273 | model_tp_dict = {} |
| 274 | if args.model_name == "": |
| 275 | model_tp_dict = { |
| 276 | "Qwen/Qwen2-VL-2B-Instruct": [1], |
| 277 | "Qwen/Qwen2-VL-7B-Instruct": [1], |
| 278 | "Qwen/Qwen2-VL-72B-Instruct": [2, 4, 8], |
| 279 | "Qwen/Qwen2.5-VL-3B-Instruct": [1, 2, 4, 8], |
| 280 | "Qwen/Qwen2.5-VL-7B-Instruct": [1, 2, 4, 8], |
| 281 | "Qwen/Qwen2.5-VL-72B-Instruct": [2, 4, 8], |
| 282 | } |
| 283 | else: |
| 284 | model_tp_dict[args.model_name] = [args.tp_size] |
| 285 | |
| 286 | if args.num_tokens is None: |
| 287 | num_tokens_list = [2**i for i in range(0, 18)] |
| 288 | else: |
| 289 | num_tokens_list = args.num_tokens |
| 290 | |
| 291 | for model_name, tp_list in model_tp_dict.items(): |
| 292 | config = get_config(model_name, trust_remote_code=args.trust_remote_code) |
| 293 | for tp_size in tp_list: |
| 294 | # get the model config |
| 295 | total_num_kv_heads = config.num_key_value_heads |
| 296 | total_num_heads = config.num_attention_heads |
| 297 | num_heads = total_num_heads // tp_size |
| 298 | num_kv_heads = max(1, total_num_kv_heads // tp_size) |
| 299 | head_dim = config.hidden_size // total_num_heads |
| 300 | q_size = num_heads * head_dim |
| 301 | kv_size = num_kv_heads * head_dim |
| 302 | is_neox_style = True |
| 303 | rope_parameters = config.rope_parameters |
| 304 | max_position = config.max_position_embeddings |
| 305 | |
| 306 | for num_tokens in num_tokens_list: |
| 307 | benchmark_mrope( |
| 308 | model_name=model_name, |
| 309 | num_tokens=num_tokens, |
| 310 | head_dim=head_dim, |
| 311 | tp_size=tp_size, |
| 312 | num_heads=num_heads, |
| 313 | num_kv_heads=num_kv_heads, |
| 314 | max_position=max_position, |
| 315 | is_neox_style=is_neox_style, |
| 316 | rope_parameters=rope_parameters, |
| 317 | dtype=getattr(torch, args.dtype), |
| 318 | seed=args.seed, |
| 319 | warmup_iter=args.warmup_iter, |
| 320 | benchmark_iter=args.benchmark_iter, |
| 321 | csv_writer=csv_writer, |
| 322 | ) |
| 323 | |
| 324 | print(f"Benchmark results saved to {csv_filename}") |