| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Comprehensive 3-way SiLU Benchmark Suite. |
| 5 | |
| 6 | This benchmark compares three SiLU implementations: |
| 7 | 1. SiLU V2 (CUDA) - Optimized CUDA kernel implementation |
| 8 | 2. Triton Kernel - Triton-based implementation |
| 9 | |
| 10 | The suite generates detailed performance comparisons including: |
| 11 | - Memory bandwidth utilization |
| 12 | - Speedup ratios (baseline vs optimized implementations) |
| 13 | - Performance across different expert configurations and token distributions |
| 14 | """ |
| 15 | |
| 16 | from collections.abc import Callable |
| 17 | |
| 18 | import matplotlib.pyplot as plt |
| 19 | import numpy as np |
| 20 | import torch |
| 21 | |
| 22 | from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( |
| 23 | persistent_masked_m_silu_mul_quant, |
| 24 | ) |
| 25 | from vllm.triton_utils import tl, triton |
| 26 | from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used |
| 27 | from vllm.utils.torch_utils import set_random_seed |
| 28 | |
| 29 | |
| 30 | @triton.jit |
| 31 | def _silu_mul_fp8_quant_deep_gemm( |
| 32 | # Pointers ------------------------------------------------------------ |
| 33 | input_ptr, # 16-bit activations (E, T, 2*H) |
| 34 | y_q_ptr, # fp8 quantized activations (E, T, H) |
| 35 | y_s_ptr, # 16-bit scales (E, T, G) |
| 36 | counts_ptr, # int32 num tokens per expert (E) |
| 37 | # Sizes --------------------------------------------------------------- |
| 38 | H: tl.constexpr, # hidden dimension (per output) |
| 39 | GROUP_SIZE: tl.constexpr, # elements per group (usually 128) |
| 40 | # Strides for input (elements) --------------------------------------- |
| 41 | stride_i_e, |
| 42 | stride_i_t, |
| 43 | stride_i_h, |
| 44 | # Strides for y_q (elements) ----------------------------------------- |
| 45 | stride_yq_e, |
| 46 | stride_yq_t, |
| 47 | stride_yq_h, |
| 48 | # Strides for y_s (elements) ----------------------------------------- |
| 49 | stride_ys_e, |
| 50 | stride_ys_t, |
| 51 | stride_ys_g, |
| 52 | # Stride for counts (elements) |
| 53 | stride_counts_e, |
| 54 | # Numeric params ------------------------------------------------------ |
| 55 | eps: tl.constexpr, |
| 56 | fp8_min: tl.constexpr, |
| 57 | fp8_max: tl.constexpr, |
| 58 | use_ue8m0: tl.constexpr, |
| 59 | # Meta --------------------------------------------------------------- |
| 60 | BLOCK: tl.constexpr, |
| 61 | NUM_STAGES: tl.constexpr, |
| 62 | ): |
| 63 | G = H // GROUP_SIZE |
| 64 | |
| 65 | # map program id -> (e, g) |
| 66 | pid = tl.program_id(0) |
| 67 | e = pid // G |
| 68 | g = pid % G |
| 69 | |
| 70 | e = e.to(tl.int64) |
| 71 | g = g.to(tl.int64) |
| 72 | |
| 73 | # number of valid tokens for this expert |
| 74 | n_tokens = tl.load(counts_ptr + e * stride_counts_e).to(tl.int64) |
| 75 | |
| 76 | cols = tl.arange(0, BLOCK).to(tl.int64) |
| 77 | mask = cols < BLOCK |
| 78 | |
| 79 | base_input_offset = e * stride_i_e + g * GROUP_SIZE * stride_i_h |
| 80 | base_gate_offset = base_input_offset + cols * stride_i_h |
| 81 | base_up_offset = base_input_offset + H * stride_i_h + cols * stride_i_h |
| 82 | base_yq_offset = e * stride_yq_e + g * GROUP_SIZE * stride_yq_h + cols * stride_yq_h |
| 83 | base_ys_offset = e * stride_ys_e + g * stride_ys_g |
| 84 | |
| 85 | for t in tl.range(0, n_tokens, num_stages=NUM_STAGES): |
| 86 | gate = tl.load( |
| 87 | input_ptr + base_gate_offset + t * stride_i_t, mask=mask, other=0.0 |
| 88 | ).to(tl.float32) |
| 89 | up = tl.load(input_ptr + base_up_offset + t * stride_i_t, mask=mask, other=0.0) |
| 90 | |
| 91 | gate = gate * (1.0 / (1.0 + tl.exp(-gate))) |
| 92 | y = gate * up |
| 93 | |
| 94 | y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max |
| 95 | if use_ue8m0: |
| 96 | y_s = tl.exp2(tl.ceil(tl.log2(y_s))) |
| 97 | |
| 98 | y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) |
| 99 | |
| 100 | tl.store(y_q_ptr + base_yq_offset + t * stride_yq_t, y_q, mask=mask) |
| 101 | tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s) |
| 102 | |
| 103 | |
| 104 | def silu_mul_fp8_quant_deep_gemm_triton( |
| 105 | y: torch.Tensor, # (E, T, 2*H) |
| 106 | tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert |
| 107 | num_parallel_tokens, |
| 108 | group_size: int = 128, |
| 109 | eps: float = 1e-10, |
| 110 | expert_offsets: torch.Tensor = None, |
| 111 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 112 | """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales. |
| 113 | |
| 114 | y has shape (E, T, 2*H). The first half of the last dimension is |
| 115 | silu-activated, multiplied by the second half, then quantized into FP8. |
| 116 | |
| 117 | Returns `(y_q, y_s)` where |
| 118 | * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H] |
| 119 | * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T) |
| 120 | """ |
| 121 | assert y.ndim == 3, "y must be (E, T, 2*H)" |
| 122 | E, T, H2 = y.shape |
| 123 | assert H2 % 2 == 0, "last dim of y must be even (2*H)" |
| 124 | H = H2 // 2 |
| 125 | G = (H + group_size - 1) // group_size |
| 126 | assert H % group_size == 0, "H must be divisible by group_size" |
| 127 | assert tokens_per_expert.ndim == 1 and tokens_per_expert.shape[0] == E, ( |
| 128 | "tokens_per_expert must be shape (E,)" |
| 129 | ) |
| 130 | tokens_per_expert = tokens_per_expert.to(device=y.device, dtype=torch.int32) |
| 131 | |
| 132 | # allocate outputs |
| 133 | fp8_dtype = torch.float8_e4m3fn |
| 134 | y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device) |
| 135 | |
| 136 | # strides (elements) |
| 137 | stride_i_e, stride_i_t, stride_i_h = y.stride() |
| 138 | stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride() |
| 139 | |
| 140 | # desired scale strides (elements): (T*G, 1, T) |
| 141 | stride_ys_e = T * G |
| 142 | stride_ys_t = 1 |
| 143 | stride_ys_g = T |
| 144 | y_s = torch.empty_strided( |
| 145 | (E, T, G), |
| 146 | (stride_ys_e, stride_ys_t, stride_ys_g), |
| 147 | dtype=torch.float32, |
| 148 | device=y.device, |
| 149 | ) |
| 150 | |
| 151 | stride_cnt_e = tokens_per_expert.stride()[0] |
| 152 | |
| 153 | # Static grid over experts and H-groups. |
| 154 | # A loop inside the kernel handles the token dim |
| 155 | grid = (E * G,) |
| 156 | |
| 157 | f_info = torch.finfo(fp8_dtype) |
| 158 | fp8_max = f_info.max |
| 159 | fp8_min = f_info.min |
| 160 | |
| 161 | _silu_mul_fp8_quant_deep_gemm[grid]( |
| 162 | y, |
| 163 | y_q, |
| 164 | y_s, |
| 165 | tokens_per_expert, |
| 166 | H, |
| 167 | group_size, |
| 168 | stride_i_e, |
| 169 | stride_i_t, |
| 170 | stride_i_h, |
| 171 | stride_yq_e, |
| 172 | stride_yq_t, |
| 173 | stride_yq_h, |
| 174 | stride_ys_e, |
| 175 | stride_ys_t, |
| 176 | stride_ys_g, |
| 177 | stride_cnt_e, |
| 178 | eps, |
| 179 | fp8_min, |
| 180 | fp8_max, |
| 181 | is_deep_gemm_e8m0_used(), |
| 182 | BLOCK=group_size, |
| 183 | NUM_STAGES=4, |
| 184 | num_warps=1, |
| 185 | ) |
| 186 | |
| 187 | return y_q, y_s |
| 188 | |
| 189 | |
| 190 | # Parse generation strategies |
| 191 | strategies = ["random_imbalanced", "uniform", "max_t"] |
| 192 | |
| 193 | |
| 194 | def benchmark( |
| 195 | kernel: Callable, |
| 196 | E: int, |
| 197 | T: int, |
| 198 | H: int, |
| 199 | total_tokens: int, |
| 200 | num_parallel_tokens: int = 64, |
| 201 | G: int = 128, |
| 202 | runs: int = 200, |
| 203 | num_warmups: int = 20, |
| 204 | gen_strategy: str = "default", |
| 205 | iterations_per_run: int = 20, |
| 206 | ): |
| 207 | def generate_data(seed_offset=0): |
| 208 | """Generate input data with given seed offset.""" |
| 209 | set_random_seed(42 + seed_offset) |
| 210 | y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous() |
| 211 | |
| 212 | if gen_strategy == "random_imbalanced": |
| 213 | |
| 214 | def generate_expert_loads(n_e, total_tokens, ratio, device="cuda"): |
| 215 | mean = total_tokens // n_e |
| 216 | min_max = mean // ratio |
| 217 | e = torch.ones(size=(E,), dtype=torch.int64, device=device) * mean |
| 218 | e[0] = min_max |
| 219 | r = torch.rand(size=(E - 1,)) |
| 220 | r /= r.sum() |
| 221 | r *= total_tokens - min_max |
| 222 | r = r.round().long() |
| 223 | e[1:] = r.to(device=device) |
| 224 | return e |
| 225 | |
| 226 | tokens_per_expert = generate_expert_loads(E, total_tokens, 0.7, "cuda") |
| 227 | elif gen_strategy == "uniform": |
| 228 | r = torch.rand(size=(E,)) |
| 229 | r /= r.sum() |
| 230 | r *= total_tokens |
| 231 | r = r.round().long() |
| 232 | tokens_per_expert = r |
| 233 | elif gen_strategy == "max_t": |
| 234 | tokens_per_expert = torch.empty(size=(E,), dtype=torch.int32, device="cuda") |
| 235 | tokens_per_expert.fill_(total_tokens / E) |
| 236 | elif gen_strategy == "first_t": |
| 237 | tokens_per_expert = torch.zeros(size=(E,), dtype=torch.int32, device="cuda") |
| 238 | tokens_per_expert[0] = min(T, total_tokens) |
| 239 | else: |
| 240 | raise ValueError(f"Unknown generation strategy: {gen_strategy}") |
| 241 | return y, tokens_per_expert |
| 242 | |
| 243 | dataset_count = 4 |
| 244 | # Pre-generate different input matrices for each iteration to avoid cache effects |
| 245 | data_sets = [generate_data(i) for i in range(dataset_count)] |
| 246 | |
| 247 | # Warmup |
| 248 | y, tokens_per_expert = data_sets[0] |
| 249 | for _ in range(num_warmups): |
| 250 | kernel( |
| 251 | y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G |
| 252 | ) |
| 253 | torch.accelerator.synchronize() |
| 254 | |
| 255 | start_event = torch.Event(enable_timing=True) |
| 256 | end_event = torch.Event(enable_timing=True) |
| 257 | |
| 258 | # Benchmark |
| 259 | latencies: list[float] = [] |
| 260 | for _ in range(runs): |
| 261 | torch.accelerator.synchronize() |
| 262 | |
| 263 | start_event.record() |
| 264 | for i in range(iterations_per_run): |
| 265 | y, tokens_per_expert = data_sets[i % dataset_count] |
| 266 | kernel( |
| 267 | y, |
| 268 | tokens_per_expert, |
| 269 | num_parallel_tokens=num_parallel_tokens, |
| 270 | group_size=G, |
| 271 | ) |
| 272 | end_event.record() |
| 273 | end_event.synchronize() |
| 274 | |
| 275 | total_time_ms = start_event.elapsed_time(end_event) |
| 276 | per_iter_time_ms = total_time_ms / iterations_per_run |
| 277 | latencies.append(per_iter_time_ms) |
| 278 | |
| 279 | # Use median instead of average for better outlier handling |
| 280 | median_time_ms = np.median(latencies) |
| 281 | median_time_s = median_time_ms / 1000 |
| 282 | |
| 283 | # Calculate actual work done (using first dataset for consistency) |
| 284 | _, tokens_per_expert = data_sets[0] |
| 285 | actual_tokens = tokens_per_expert.sum().item() |
| 286 | actual_elements = actual_tokens * H |
| 287 | |
| 288 | # GFLOPS: operations per element = exp + 3 muls + 1 div + quantization ops ≈ 8 ops |
| 289 | ops_per_element = 8 |
| 290 | total_ops = actual_elements * ops_per_element |
| 291 | gflops = total_ops / median_time_s / 1e9 |
| 292 | |
| 293 | # Memory bandwidth: bfloat16 inputs (2 bytes), fp8 output (1 byte), scales (4 bytes) |
| 294 | input_bytes = actual_tokens * 2 * H * 2 # 2*H bfloat16 inputs |
| 295 | output_bytes = actual_tokens * H * 1 # H fp8 outputs |
| 296 | scale_bytes = actual_tokens * (H // G) * 4 # scales in float32 |
| 297 | total_bytes = input_bytes + output_bytes + scale_bytes |
| 298 | memory_bw = total_bytes / median_time_s / 1e9 |
| 299 | |
| 300 | HOPPER_BANDWIDTH_TBPS = 3.35 |
| 301 | return ( |
| 302 | median_time_ms, |
| 303 | gflops, |
| 304 | memory_bw, |
| 305 | (memory_bw / (HOPPER_BANDWIDTH_TBPS * 1024)) * 100, |
| 306 | ) |
| 307 | |
| 308 | |
| 309 | def create_comparison_plot( |
| 310 | ratios, silu_v2_times, triton_times, config_labels, strategy_name, id |
| 311 | ): |
| 312 | fig, ax = plt.subplots(1, 1, figsize=(18, 6)) |
| 313 | |
| 314 | # Configure x-axis positions |
| 315 | x = np.arange(len(config_labels)) |
| 316 | width = 0.25 |
| 317 | |
| 318 | # Execution Time plot (lower is better) |
| 319 | ax.bar(x, silu_v2_times, width, label="SiLU V2 (CUDA)", alpha=0.8, color="blue") |
| 320 | ax.bar( |
| 321 | x + width, triton_times, width, label="Triton Kernel", alpha=0.8, color="green" |
| 322 | ) |
| 323 | |
| 324 | # Add speedup labels over each bar trio |
| 325 | for i in range(len(x)): |
| 326 | triton_v2_speedup = ratios[i][1] # triton/v2 |
| 327 | max_height = max(silu_v2_times[i], triton_times[i]) |
| 328 | |
| 329 | # Triton/V2 speedup |
| 330 | ax.text( |
| 331 | x[i] + width / 2, |
| 332 | max_height + max_height * 0.02, |
| 333 | f"{triton_v2_speedup:.2f}x", |
| 334 | ha="center", |
| 335 | va="bottom", |
| 336 | fontweight="bold", |
| 337 | fontsize=8, |
| 338 | ) |
| 339 | |
| 340 | ax.set_xlabel("Configuration") |
| 341 | ax.set_ylabel("% Utilization") |
| 342 | ax.set_title( |
| 343 | f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)" |
| 344 | ) |
| 345 | ax.set_xticks(x) |
| 346 | ax.set_xticklabels(config_labels, rotation=45, ha="right") |
| 347 | ax.legend() |
| 348 | ax.grid(True, alpha=0.3) |
| 349 | |
| 350 | plt.tight_layout() |
| 351 | return fig, ax |
| 352 | |
| 353 | |
| 354 | def create_combined_plot(all_results): |
| 355 | num_strategies = len(all_results) |
| 356 | fig, axes = plt.subplots(num_strategies, 1, figsize=(22, 7 * num_strategies)) |
| 357 | |
| 358 | if num_strategies == 1: |
| 359 | axes = [axes] |
| 360 | |
| 361 | for idx, ( |
| 362 | strategy_name, |
| 363 | all_ratios, |
| 364 | all_silu_v2_results, |
| 365 | all_triton_results, |
| 366 | config_labels, |
| 367 | config_x_axis, |
| 368 | ) in enumerate(all_results): |
| 369 | ax = axes[idx] |
| 370 | |
| 371 | # Flatten the nested results to get bandwidth percentages for plotting |
| 372 | silu_v2_bandwidths = [] |
| 373 | triton_bandwidths = [] |
| 374 | flat_ratios = [] |
| 375 | |
| 376 | for config_results in all_silu_v2_results: |
| 377 | for result in config_results: |
| 378 | silu_v2_bandwidths.append(result[3]) # bandwidth percentage |
| 379 | |
| 380 | for config_results in all_triton_results: |
| 381 | for result in config_results: |
| 382 | triton_bandwidths.append(result[3]) # bandwidth percentage |
| 383 | |
| 384 | for config_ratios in all_ratios: |
| 385 | for ratio in config_ratios: |
| 386 | flat_ratios.append(ratio) |
| 387 | |
| 388 | # Configure x-axis positions |
| 389 | x = np.arange(len(config_labels)) |
| 390 | width = 0.25 |
| 391 | |
| 392 | # Bandwidth utilization plot (higher is better) |
| 393 | ax.bar( |
| 394 | x, |
| 395 | silu_v2_bandwidths, |
| 396 | width, |
| 397 | label="SiLU V2 (CUDA)", |
| 398 | alpha=0.8, |
| 399 | color="blue", |
| 400 | ) |
| 401 | ax.bar( |
| 402 | x + width, |
| 403 | triton_bandwidths, |
| 404 | width, |
| 405 | label="Triton Kernel", |
| 406 | alpha=0.8, |
| 407 | color="green", |
| 408 | ) |
| 409 | |
| 410 | # Add speedup labels over each bar trio |
| 411 | for i in range(len(x)): |
| 412 | triton_v2_speedup = flat_ratios[i] # triton/v2 |
| 413 | max_height = max(silu_v2_bandwidths[i], triton_bandwidths[i]) |
| 414 | |
| 415 | # Triton/V2 speedup |
| 416 | ax.text( |
| 417 | x[i] + width / 2, |
| 418 | max_height + max_height * 0.02, |
| 419 | f"{triton_v2_speedup:.2f}x", |
| 420 | ha="center", |
| 421 | va="bottom", |
| 422 | fontweight="bold", |
| 423 | fontsize=8, |
| 424 | ) |
| 425 | |
| 426 | ax.set_xlabel("Configuration") |
| 427 | ax.set_ylabel("% Utilization") |
| 428 | ax.set_title( |
| 429 | f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)" |
| 430 | ) |
| 431 | ax.set_xticks(x) |
| 432 | ax.set_xticklabels(config_labels, rotation=45, ha="right") |
| 433 | ax.legend() |
| 434 | ax.grid(True, alpha=0.3) |
| 435 | |
| 436 | plt.tight_layout() |
| 437 | filename = "silu_benchmark_combined_3way.png" |
| 438 | plt.savefig(filename, dpi=300, bbox_inches="tight") |
| 439 | plt.show() |
| 440 | |
| 441 | return filename |
| 442 | |
| 443 | |
| 444 | outer_dim = 7168 |
| 445 | configs = [ |
| 446 | # DeepSeekV3 Configs |
| 447 | # (1, 56, 7168), |
| 448 | (8, 1024, 7168), |
| 449 | # (32, 56, 7168), |
| 450 | # DeepSeekV3 Configs |
| 451 | (32, 1024, 7168), |
| 452 | # DeepSeekV3 Configs |
| 453 | (256, 1024, 7168), |
| 454 | ] |
| 455 | |
| 456 | runs = 100 |
| 457 | num_warmups = 20 |
| 458 | |
| 459 | strategy_descriptions = { |
| 460 | "uniform": "Uniform Random", |
| 461 | "random_imbalanced": "Imbalanced Random", |
| 462 | "max_t": "Even Assignment", |
| 463 | "first_t": "experts[0] = T, experts[1:] = 0", |
| 464 | } |
| 465 | |
| 466 | print(f"GPU: {torch.cuda.get_device_name()}") |
| 467 | print(f"Testing strategies: {', '.join(strategies)}") |
| 468 | print(f"Configurations: {len(configs)} configs") |
| 469 | |
| 470 | all_results = [] |
| 471 | |
| 472 | # Run benchmarks for each strategy |
| 473 | for id, strategy in enumerate(strategies): |
| 474 | print(f"\n{'=' * 60}") |
| 475 | print(f"Testing strategy: {strategy_descriptions[strategy]}") |
| 476 | print(f"{'=' * 60}") |
| 477 | |
| 478 | # Collect benchmark data for all three algorithms |
| 479 | config_labels = [] |
| 480 | config_x_axis = [] |
| 481 | all_silu_v2_results = [] |
| 482 | all_triton_results = [] |
| 483 | all_ratios = [] |
| 484 | |
| 485 | for E, T, H in configs: |
| 486 | total_tokens_config = [] |
| 487 | for i in [8, 16, 32, 64, 128, 256, 512]: |
| 488 | if i <= T: |
| 489 | total_tokens_config.append(i * E) |
| 490 | config_x_axis.append(total_tokens_config) |
| 491 | |
| 492 | silu_v2_results = [] |
| 493 | triton_results = [] |
| 494 | ratios = [] |
| 495 | |
| 496 | for total_tokens in total_tokens_config: |
| 497 | config_label = f"E={E},T={T},H={H},TT={total_tokens}" |
| 498 | config_labels.append(config_label) |
| 499 | |
| 500 | # SiLU V2 (CUDA kernel) results |
| 501 | time_ms_silu_v2, gflops, gbps, perc = benchmark( |
| 502 | persistent_masked_m_silu_mul_quant, |
| 503 | E, |
| 504 | T, |
| 505 | H, |
| 506 | total_tokens, |
| 507 | runs=runs, |
| 508 | num_warmups=num_warmups, |
| 509 | gen_strategy=strategy, |
| 510 | ) |
| 511 | silu_v2_results.append((time_ms_silu_v2, gflops, gbps, perc)) |
| 512 | |
| 513 | # Triton kernel results |
| 514 | time_ms_triton, gflops, gbps, perc = benchmark( |
| 515 | silu_mul_fp8_quant_deep_gemm_triton, |
| 516 | E, |
| 517 | T, |
| 518 | H, |
| 519 | total_tokens, |
| 520 | runs=runs, |
| 521 | num_warmups=num_warmups, |
| 522 | gen_strategy=strategy, |
| 523 | ) |
| 524 | triton_results.append((time_ms_triton, gflops, gbps, perc)) |
| 525 | |
| 526 | # Calculate speedup ratios (triton baseline / implementation) |
| 527 | triton_v2_ratio = time_ms_triton / time_ms_silu_v2 |
| 528 | ratios.append(triton_v2_ratio) |
| 529 | |
| 530 | print( |
| 531 | f"Completed: {config_label}:" |
| 532 | f" V2: {time_ms_silu_v2:.3f}ms," |
| 533 | f" Triton: {time_ms_triton:.3f}ms" |
| 534 | ) |
| 535 | |
| 536 | all_silu_v2_results.append(silu_v2_results) |
| 537 | all_triton_results.append(triton_results) |
| 538 | all_ratios.append(ratios) |
| 539 | |
| 540 | # Store results for combined plotting |
| 541 | all_results.append( |
| 542 | ( |
| 543 | strategy_descriptions[strategy], |
| 544 | all_ratios, |
| 545 | all_silu_v2_results, |
| 546 | all_triton_results, |
| 547 | config_labels, |
| 548 | config_x_axis, |
| 549 | ) |
| 550 | ) |
| 551 | |
| 552 | # Print summary table for this strategy |
| 553 | print(f"\nSummary Table - {strategy_descriptions[strategy]}:") |
| 554 | print(f" {'V2 Time(ms)':<12} {'Triton Time(ms)':<14} {'Triton/V2':<10}") |
| 555 | print("-" * 90) |
| 556 | |
| 557 | for i, (E, T, H) in enumerate(configs): |
| 558 | # Get the first result for each config (simplifying for summary) |
| 559 | v2_time = silu_v2_results[i][0] |
| 560 | triton_time = triton_results[i][0] |
| 561 | triton_v2_speedup = triton_time / v2_time |
| 562 | config_label = f"E={E:3d},T={T:4d},H={H:4d}" |
| 563 | print( |
| 564 | f"{config_label:<20} {v2_time:8.5f} {triton_time:10.5f} " |
| 565 | f"{triton_v2_speedup:8.2f}x" |
| 566 | ) |
| 567 | |
| 568 | |
| 569 | def create_total_tokens_plot(all_results): |
| 570 | num_strategies = len(all_results) |
| 571 | num_configs = len(configs) |
| 572 | |
| 573 | fig, axs = plt.subplots( |
| 574 | num_strategies, num_configs * 2, figsize=(32, 8 * num_strategies) |
| 575 | ) |
| 576 | |
| 577 | # Add main title to the entire figure |
| 578 | fig.suptitle( |
| 579 | "Performance Analysis: Speedup vs Bandwidth Utilization (SiLU V2, and Triton)", |
| 580 | fontsize=18, |
| 581 | fontweight="bold", |
| 582 | y=0.98, |
| 583 | ) |
| 584 | |
| 585 | # Handle single strategy case |
| 586 | if num_strategies == 1: |
| 587 | axs = axs.reshape(1, -1) |
| 588 | |
| 589 | # Handle single config case |
| 590 | if num_configs == 1: |
| 591 | axs = axs.reshape(-1, 2) |
| 592 | |
| 593 | for strategy_idx, result in enumerate(all_results): |
| 594 | ( |
| 595 | strategy_name, |
| 596 | all_ratios, |
| 597 | all_silu_v2_results, |
| 598 | all_triton_results, |
| 599 | config_labels, |
| 600 | config_x_axis, |
| 601 | ) = result |
| 602 | |
| 603 | for config_idx in range(num_configs): |
| 604 | # Speedup plot (left column) |
| 605 | ax_speedup = axs[strategy_idx, config_idx * 2] |
| 606 | # Bandwidth plot (right column) |
| 607 | ax_bandwidth = axs[strategy_idx, config_idx * 2 + 1] |
| 608 | |
| 609 | E, T, H = configs[config_idx] |
| 610 | ratios = all_ratios[config_idx] |
| 611 | total_tokens_values = config_x_axis[config_idx] |
| 612 | |
| 613 | # Extract speedup ratios |
| 614 | triton_v2_ratios = [ratio for ratio in ratios] |
| 615 | |
| 616 | # Extract bandwidth percentages for all implementations |
| 617 | v2_bandwidth_percentages = [ |
| 618 | result[3] for result in all_silu_v2_results[config_idx] |
| 619 | ] |
| 620 | triton_bandwidth_percentages = [ |
| 621 | result[3] for result in all_triton_results[config_idx] |
| 622 | ] |
| 623 | |
| 624 | # Plot speedup ratios vs total tokens (left plot) |
| 625 | ax_speedup.plot( |
| 626 | total_tokens_values, |
| 627 | triton_v2_ratios, |
| 628 | "go-", |
| 629 | linewidth=3, |
| 630 | markersize=8, |
| 631 | label="Triton/V2 Speedup", |
| 632 | ) |
| 633 | ax_speedup.set_title( |
| 634 | f"{strategy_name}\nSpeedup vs Baseline (Triton)\nE={E}, T={T}, H={H}", |
| 635 | fontsize=12, |
| 636 | fontweight="bold", |
| 637 | ) |
| 638 | ax_speedup.set_xlabel("Total Tokens", fontweight="bold", fontsize=11) |
| 639 | ax_speedup.set_ylabel("Speedup Ratio", fontweight="bold", fontsize=11) |
| 640 | ax_speedup.legend(prop={"weight": "bold"}) |
| 641 | ax_speedup.grid(True, alpha=0.3) |
| 642 | |
| 643 | # Plot bandwidth utilization (right plot) |
| 644 | ax_bandwidth.plot( |
| 645 | total_tokens_values, |
| 646 | v2_bandwidth_percentages, |
| 647 | "o-", |
| 648 | linewidth=3, |
| 649 | markersize=8, |
| 650 | label="SiLU V2", |
| 651 | color="blue", |
| 652 | ) |
| 653 | ax_bandwidth.plot( |
| 654 | total_tokens_values, |
| 655 | triton_bandwidth_percentages, |
| 656 | "o-", |
| 657 | linewidth=3, |
| 658 | markersize=8, |
| 659 | label="Triton", |
| 660 | color="green", |
| 661 | ) |
| 662 | ax_bandwidth.set_title( |
| 663 | f"{strategy_name}\nBandwidth Utilization (Hopper)\nE={E}, T={T}, H={H}", |
| 664 | fontsize=12, |
| 665 | fontweight="bold", |
| 666 | ) |
| 667 | ax_bandwidth.set_xlabel("Total Tokens", fontweight="bold", fontsize=11) |
| 668 | ax_bandwidth.set_ylabel( |
| 669 | "% of Peak Bandwidth", fontweight="bold", fontsize=11 |
| 670 | ) |
| 671 | ax_bandwidth.legend(prop={"weight": "bold"}) |
| 672 | ax_bandwidth.grid(True, alpha=0.3) |
| 673 | |
| 674 | # Format x-axis labels for both plots |
| 675 | for ax in [ax_speedup, ax_bandwidth]: |
| 676 | ax.set_xticks(total_tokens_values) |
| 677 | ax.set_xticklabels( |
| 678 | [ |
| 679 | f"{tt // 1000}K" if tt >= 1000 else str(tt) |
| 680 | for tt in total_tokens_values |
| 681 | ], |
| 682 | fontweight="bold", |
| 683 | ) |
| 684 | # Make tick labels bold |
| 685 | for label in ax.get_xticklabels() + ax.get_yticklabels(): |
| 686 | label.set_fontweight("bold") |
| 687 | |
| 688 | # Add value labels on Triton/V2 speedup points |
| 689 | for x, y in zip(total_tokens_values, triton_v2_ratios): |
| 690 | ax_speedup.annotate( |
| 691 | f"{y:.2f}x", |
| 692 | (x, y), |
| 693 | textcoords="offset points", |
| 694 | xytext=(0, -15), |
| 695 | ha="center", |
| 696 | fontsize=9, |
| 697 | fontweight="bold", |
| 698 | bbox=dict(boxstyle="round,pad=0.2", facecolor="green", alpha=0.3), |
| 699 | ) |
| 700 | |
| 701 | plt.tight_layout() |
| 702 | plt.subplots_adjust(top=0.93) # Make room for main title |
| 703 | filename = "silu_benchmark_total_tokens_3way.png" |
| 704 | plt.savefig(filename, dpi=300, bbox_inches="tight") |
| 705 | plt.show() |
| 706 | |
| 707 | return filename |
| 708 | |
| 709 | |
| 710 | # Create comprehensive 3-way comparison plots |
| 711 | combined_plot_filename = create_combined_plot(all_results) |
| 712 | total_tokens_plot_filename = create_total_tokens_plot(all_results) |
| 713 | |
| 714 | print(f"\n{'=' * 80}") |
| 715 | print("3-Way Benchmark Suite Complete!") |
| 716 | print(f"Generated combined comparison plot: {combined_plot_filename}") |
| 717 | print(f"Generated total tokens analysis plot: {total_tokens_plot_filename}") |
| 718 | print("Compared: SiLU V2 (CUDA), and Triton implementations") |
| 719 | print(f"{'=' * 80}") |