@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ca100c90fe
Fork 0

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 15, 21:00 UTC

Add benchmark for DeepGEMM and vLLM Block FP8 Dense GEMM (#13917)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Mar 5, 2025 at 20:08 UTC ca100c90fef6cf094c856926da9a6abd855743fa
2 files changed +593
benchmarks/kernels/deepgemm/README.md new
+129
@@ -0,0 +1,129 @@
1 +# DeepSeek DeepGEMM Kernels Benchmark
2 +
3 +This directory includes benchmarks between DeepSeek's DeepGEMM block fp8 kernels against vLLM's existing triton and CUTLASS-based kernels.
4 +
5 +Currently this just includes dense GEMMs and only works on Hopper GPUs.
6 +
7 +## Setup
8 +
9 +You need to install vLLM in your usual fashion, then install DeepGEMM from source in its own directory:
10 +
11 +```
12 +git clone --recursive https://github.com/deepseek-ai/DeepGEMM
13 +cd DeepGEMM
14 +python setup.py install
15 +uv pip install -e .
16 +```
17 +
18 +## Usage
19 +
20 +```
21 +python benchmark_fp8_block_dense_gemm.py
22 +INFO 02-26 21:55:13 [__init__.py:207] Automatically detected platform cuda.
23 +===== STARTING FP8 GEMM BENCHMARK =====
24 +PyTorch version: 2.5.1+cu124
25 +CUDA version: 12.4
26 +Triton version: 3.1.0
27 +Using device: NVIDIA H100 80GB HBM3
28 +WARNING 02-26 21:55:15 [fp8_utils.py:458] Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! Config file not found at /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
29 +INFO 02-26 21:55:15 [fp8_utils.py:449] Using configuration from /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=18432,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json for W8A8 Block FP8 kernel.
30 +WARNING 02-26 21:55:16 [fp8_utils.py:458] Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! Config file not found at /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=18432,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
31 +WARNING 02-26 21:55:17 [fp8_utils.py:458] Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! Config file not found at /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=1536,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
32 +INFO 02-26 21:55:17 [fp8_utils.py:449] Using configuration from /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=32768,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json for W8A8 Block FP8 kernel.
33 +INFO 02-26 21:55:17 [fp8_utils.py:449] Using configuration from /home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json for W8A8 Block FP8 kernel.
34 +
35 +===== PERFORMANCE COMPARISON =====
36 +
37 +DeepGEMM Implementation:
38 ++------+-------+-------+-----------+--------+--------+
39 +| m | n | k | Time (μs) | TFLOPS | GB/s |
40 ++------+-------+-------+-----------+--------+--------+
41 +| 8 | 4096 | 7168 | 102.9 | 4.6 | 286.4 |
42 +| 8 | 7168 | 18432 | 70.8 | 29.8 | 1868.8 |
43 +| 8 | 18432 | 7168 | 69.3 | 30.5 | 1911.8 |
44 +| 64 | 4096 | 7168 | 69.1 | 54.4 | 439.0 |
45 +| 64 | 7168 | 18432 | 69.4 | 243.6 | 1933.6 |
46 +| 64 | 18432 | 7168 | 70.4 | 240.3 | 1917.2 |
47 +| 64 | 24576 | 1536 | 70.1 | 68.9 | 584.6 |
48 +| 64 | 32768 | 512 | 68.4 | 31.4 | 307.1 |
49 +| 64 | 7168 | 16384 | 69.5 | 216.3 | 1718.5 |
50 +| 128 | 4096 | 7168 | 141.1 | 53.3 | 222.1 |
51 +| 128 | 7168 | 18432 | 71.9 | 470.5 | 1896.1 |
52 +| 128 | 18432 | 7168 | 69.3 | 488.2 | 1988.2 |
53 +| 1024 | 4096 | 7168 | 89.7 | 670.1 | 502.5 |
54 +| 1024 | 18432 | 7168 | 279.0 | 969.8 | 635.2 |
55 +| 2048 | 4096 | 7168 | 175.1 | 687.0 | 347.4 |
56 +| 4096 | 4096 | 7168 | 335.4 | 717.0 | 275.1 |
57 ++------+-------+-------+-----------+--------+--------+
58 +
59 +vLLM Triton Implementation:
60 ++------+-------+-------+-----------+--------+--------+--------------+
61 +| m | n | k | Time (μs) | TFLOPS | GB/s | vs DeepGEMM |
62 ++------+-------+-------+-----------+--------+--------+--------------+
63 +| 8 | 4096 | 7168 | 74.0 | 6.3 | 398.2 | 1.39x faster |
64 +| 8 | 7168 | 18432 | 89.6 | 23.6 | 1478.1 | 0.79x slower |
65 +| 8 | 18432 | 7168 | 113.2 | 18.7 | 1170.4 | 0.61x slower |
66 +| 64 | 4096 | 7168 | 79.4 | 47.3 | 382.2 | 0.87x slower |
67 +| 64 | 7168 | 18432 | 98.5 | 171.7 | 1363.0 | 0.70x slower |
68 +| 64 | 18432 | 7168 | 119.5 | 141.5 | 1129.4 | 0.59x slower |
69 +| 64 | 24576 | 1536 | 37.6 | 128.4 | 1089.7 | 1.86x faster |
70 +| 64 | 32768 | 512 | 38.7 | 55.5 | 542.6 | 1.77x faster |
71 +| 64 | 7168 | 16384 | 86.1 | 174.5 | 1386.4 | 0.81x slower |
72 +| 128 | 4096 | 7168 | 90.7 | 82.9 | 345.4 | 1.56x faster |
73 +| 128 | 7168 | 18432 | 144.0 | 234.9 | 946.9 | 0.50x slower |
74 +| 128 | 18432 | 7168 | 229.5 | 147.4 | 600.1 | 0.30x slower |
75 +| 1024 | 4096 | 7168 | 242.3 | 248.2 | 186.1 | 0.37x slower |
76 +| 1024 | 18432 | 7168 | 897.8 | 301.4 | 197.4 | 0.31x slower |
77 +| 2048 | 4096 | 7168 | 463.0 | 259.7 | 131.4 | 0.38x slower |
78 +| 4096 | 4096 | 7168 | 901.8 | 266.7 | 102.3 | 0.37x slower |
79 ++------+-------+-------+-----------+--------+--------+--------------+
80 +
81 +vLLM CUTLASS Implementation:
82 ++------+-------+-------+-----------+--------+--------+--------------+--------------+
83 +| m | n | k | Time (μs) | TFLOPS | GB/s | vs DeepGEMM | vs Triton |
84 ++------+-------+-------+-----------+--------+--------+--------------+--------------+
85 +| 8 | 4096 | 7168 | 34.6 | 13.6 | 852.3 | 2.98x faster | 2.14x faster |
86 +| 8 | 7168 | 18432 | 78.9 | 26.8 | 1677.3 | 0.90x slower | 1.13x faster |
87 +| 8 | 18432 | 7168 | 81.2 | 26.0 | 1631.1 | 0.85x slower | 1.39x faster |
88 +| 64 | 4096 | 7168 | 36.9 | 101.9 | 822.9 | 1.87x faster | 2.15x faster |
89 +| 64 | 7168 | 18432 | 87.4 | 193.4 | 1535.2 | 0.79x slower | 1.13x faster |
90 +| 64 | 18432 | 7168 | 85.0 | 199.0 | 1587.6 | 0.83x slower | 1.41x faster |
91 +| 64 | 24576 | 1536 | 28.0 | 172.8 | 1465.8 | 2.51x faster | 1.35x faster |
92 +| 64 | 32768 | 512 | 28.8 | 74.5 | 728.5 | 2.37x faster | 1.34x faster |
93 +| 64 | 7168 | 16384 | 77.9 | 193.0 | 1532.8 | 0.89x slower | 1.11x faster |
94 +| 128 | 4096 | 7168 | 39.1 | 192.4 | 802.0 | 3.61x faster | 2.32x faster |
95 +| 128 | 7168 | 18432 | 93.7 | 360.8 | 1454.2 | 0.77x slower | 1.54x faster |
96 +| 128 | 18432 | 7168 | 85.7 | 394.8 | 1608.0 | 0.81x slower | 2.68x faster |
97 +| 1024 | 4096 | 7168 | 99.7 | 603.1 | 452.2 | 0.90x slower | 2.43x faster |
98 +| 1024 | 18432 | 7168 | 331.3 | 816.7 | 534.9 | 0.84x slower | 2.71x faster |
99 +| 2048 | 4096 | 7168 | 198.3 | 606.6 | 306.7 | 0.88x slower | 2.34x faster |
100 +| 4096 | 4096 | 7168 | 392.2 | 613.2 | 235.3 | 0.86x slower | 2.30x faster |
101 ++------+-------+-------+-----------+--------+--------+--------------+--------------+
102 +
103 +===== AVERAGE PERFORMANCE =====
104 ++----------------+------------+----------+---------------+
105 +| Implementation | Avg TFLOPS | Avg GB/s | Avg Time (ms) |
106 ++----------------+------------+----------+---------------+
107 +| DeepGEMM | 310.98 | 1052.10 | 0.11 |
108 +| vLLM Triton | 144.30 | 715.60 | 0.23 |
109 +| vLLM CUTLASS | 286.78 | 1076.67 | 0.11 |
110 ++----------------+------------+----------+---------------+
111 +
112 +===== AVERAGE SPEEDUPS =====
113 ++-----------------------------+--------------+
114 +| Comparison | Speedup |
115 ++-----------------------------+--------------+
116 +| DeepGEMM vs vLLM Triton | 1.71x faster |
117 +| DeepGEMM vs vLLM CUTLASS | 0.94x slower |
118 +| vLLM CUTLASS vs vLLM Triton | 1.84x faster |
119 ++-----------------------------+--------------+
120 +
121 +===== ACCURACY COMPARISON =====
122 ++----------------+-----------------------+
123 +| Implementation | Avg Diff vs Reference |
124 ++----------------+-----------------------+
125 +| DeepGEMM | 0.000684 |
126 +| vLLM Triton | 0.000684 |
127 +| vLLM CUTLASS | 0.000684 |
128 ++----------------+-----------------------+
129 +```
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py new
+464
@@ -0,0 +1,464 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# fmt: off
3 +# ruff: noqa: E501
4 +import time
5 +
6 +# Import DeepGEMM functions
7 +import deep_gemm
8 +import torch
9 +import triton
10 +from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
11 +
12 +# Import vLLM functions
13 +from vllm import _custom_ops as ops
14 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
15 + per_token_group_quant_fp8, w8a8_block_fp8_matmul)
16 +
17 +
18 +# Copied from
19 +# https://github.com/deepseek-ai/DeepGEMM/blob/78cacf70d41d15d688bd493ebc85845f7f2a3d5d/tests/test_core.py#L9
20 +def per_token_cast_to_fp8(
21 + x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
22 + """Convert tensor to FP8 format with per-token scaling."""
23 + assert x.dim() == 2 and x.size(1) % 128 == 0
24 + m, n = x.shape
25 + x_view = x.view(m, -1, 128)
26 + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
27 + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(
28 + torch.float8_e4m3fn).view(m, n), (x_amax / 448.0).view(m, -1)
29 +
30 +
31 +# Copied from
32 +# https://github.com/deepseek-ai/DeepGEMM/blob/78cacf70d41d15d688bd493ebc85845f7f2a3d5d/tests/test_core.py#L17
33 +def per_block_cast_to_fp8(
34 + x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
35 + """Convert tensor to FP8 format with per-block scaling."""
36 + assert x.dim() == 2
37 + m, n = x.shape
38 + x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128),
39 + dtype=x.dtype,
40 + device=x.device)
41 + x_padded[:m, :n] = x
42 + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
43 + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
44 + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
45 + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (
46 + x_amax / 448.0).view(x_view.size(0), x_view.size(2))
47 +
48 +
49 +def benchmark_shape(m: int,
50 + n: int,
51 + k: int,
52 + warmup: int = 100,
53 + repeat: int = 10000,
54 + verbose: bool = False) -> dict:
55 + """Benchmark all implementations for a specific (m, n, k) shape."""
56 + if verbose:
57 + print(f"\n=== Benchmarking shape: m={m}, n={n}, k={k} ===")
58 +
59 + # Create test tensors
60 + A = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
61 + B = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
62 +
63 + # Reference result in BF16
64 + torch.cuda.synchronize()
65 + C_ref = A @ B.t()
66 +
67 + # Pre-quantize B for all implementations
68 + # (weights can be pre-quantized offline)
69 + B_deepgemm, B_scale_deepgemm = per_block_cast_to_fp8(B)
70 + B_vllm, B_scale_vllm = per_block_cast_to_fp8(B)
71 +
72 + # Block size configuration
73 + block_size = [128, 128]
74 +
75 + # Pre-quantize A for all implementations
76 + A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)
77 + A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
78 + C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
79 + A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
80 + A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
81 + A, block_size[1], column_major_scales=True)
82 +
83 + # === DeepGEMM Implementation ===
84 + def deepgemm_gemm():
85 + # A quantization is inside the loop as it depends on activations
86 + # A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)
87 + # A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(
88 + # A, block_size[1])
89 + # A_scale_aligned = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
90 + # C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
91 + deep_gemm.gemm_fp8_fp8_bf16_nt((A_deepgemm, A_scale_deepgemm),
92 + (B_deepgemm, B_scale_deepgemm),
93 + C_deepgemm)
94 + return C_deepgemm
95 +
96 + # === vLLM Triton Implementation ===
97 + def vllm_triton_gemm():
98 + # A quantization is inside the loop as it depends on activations
99 + # A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
100 + return w8a8_block_fp8_matmul(A_vllm,
101 + B_vllm,
102 + A_scale_vllm,
103 + B_scale_vllm,
104 + block_size,
105 + output_dtype=torch.bfloat16)
106 +
107 + # === vLLM CUTLASS Implementation ===
108 + def vllm_cutlass_gemm():
109 + # A quantization is inside the loop as it depends on activations
110 + # A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
111 + # A, block_size[1], column_major_scales=True)
112 + return ops.cutlass_scaled_mm(A_vllm_cutlass,
113 + B_vllm.T,
114 + scale_a=A_scale_vllm_cutlass,
115 + scale_b=B_scale_vllm.T,
116 + out_dtype=torch.bfloat16)
117 +
118 + # Run correctness check first
119 + if verbose:
120 + print("Running correctness check...")
121 + C_deepgemm = deepgemm_gemm()
122 + C_vllm_triton = vllm_triton_gemm()
123 + C_vllm_cutlass = vllm_cutlass_gemm()
124 +
125 + deepgemm_diff = calc_diff(C_deepgemm, C_ref)
126 + vllm_triton_diff = calc_diff(C_vllm_triton, C_ref)
127 + vllm_cutlass_diff = calc_diff(C_vllm_cutlass, C_ref)
128 +
129 + if verbose:
130 + print(f"DeepGEMM vs Reference difference: {deepgemm_diff:.6f}")
131 + print(f"vLLM Triton vs Reference difference: {vllm_triton_diff:.6f}")
132 + print(f"vLLM CUTLASS vs Reference difference: {vllm_cutlass_diff:.6f}")
133 + print("vLLM Triton vs DeepGEMM difference: "
134 + f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}")
135 + print("vLLM CUTLASS vs DeepGEMM difference: "
136 + f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}")
137 +
138 + # Benchmark implementations
139 + implementations = {
140 + "DeepGEMM": deepgemm_gemm,
141 + "vLLM Triton": vllm_triton_gemm,
142 + "vLLM CUTLASS": vllm_cutlass_gemm
143 + }
144 +
145 + benchmark_results = {
146 + "shape": {
147 + "m": m,
148 + "n": n,
149 + "k": k
150 + },
151 + "implementations": {}
152 + }
153 +
154 + for name, func in implementations.items():
155 + # Warmup
156 + for _ in range(warmup):
157 + func()
158 + torch.cuda.synchronize()
159 +
160 + # Timing loop
161 + torch.cuda.synchronize()
162 + start = time.time()
163 + for _ in range(repeat):
164 + func()
165 + torch.cuda.synchronize()
166 + end = time.time()
167 +
168 + # Calculate timing and TFLOPS
169 + avg_time_ms = (end - start) / repeat * 1000
170 + avg_time_us = avg_time_ms * 1000
171 + tflops = 2 * m * n * k / (avg_time_ms * 1e-3) / 1e12
172 + gb_s = (m * k + k * n + m * n * 2) / 1e9 / (avg_time_ms * 1e-3)
173 +
174 + benchmark_results["implementations"][name] = {
175 + "time_ms": avg_time_ms,
176 + "time_us": avg_time_us,
177 + "tflops": tflops,
178 + "gb_s": gb_s,
179 + "diff": {
180 + "DeepGEMM":
181 + 0.0 if name == "DeepGEMM" else calc_diff(func(), C_deepgemm),
182 + "Reference":
183 + deepgemm_diff if name == "DeepGEMM" else
184 + (vllm_triton_diff
185 + if name == "vLLM Triton" else vllm_cutlass_diff)
186 + }
187 + }
188 +
189 + if verbose:
190 + print(
191 + f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s"
192 + )
193 +
194 + # Calculate speedups
195 + baseline = benchmark_results["implementations"]["DeepGEMM"]["time_ms"]
196 + for name, data in benchmark_results["implementations"].items():
197 + if name != "DeepGEMM":
198 + speedup = baseline / data["time_ms"]
199 + benchmark_results["implementations"][name][
200 + "speedup_vs_deepgemm"] = speedup
201 + if verbose:
202 + print(f"DeepGEMM is {1/speedup:.2f}x "
203 + f"{'faster' if 1/speedup > 1 else 'slower'} than {name}")
204 +
205 + vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"][
206 + "time_ms"]
207 + vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"][
208 + "time_ms"]
209 + cutlass_vs_triton = vllm_triton_time / vllm_cutlass_time
210 + benchmark_results["implementations"]["vLLM CUTLASS"][
211 + "speedup_vs_triton"] = cutlass_vs_triton
212 + if verbose:
213 + print(
214 + f"vLLM CUTLASS is {cutlass_vs_triton:.2f}x "
215 + f"{'faster' if cutlass_vs_triton > 1 else 'slower'} than vLLM Triton"
216 + )
217 +
218 + return benchmark_results
219 +
220 +
221 +def format_table_row(values, widths):
222 + """Format a row with specified column widths."""
223 + return "| " + " | ".join(f"{val:{w}}"
224 + for val, w in zip(values, widths)) + " |"
225 +
226 +
227 +def print_table(headers, rows, title=None):
228 + """Print a table with headers and rows."""
229 + if title:
230 + print(f"\n{title}")
231 +
232 + # Calculate column widths based on headers and data
233 + widths = [
234 + max(len(str(h)), max(len(str(row[i])) for row in rows))
235 + for i, h in enumerate(headers)
236 + ]
237 +
238 + # Create separator line
239 + separator = "+-" + "-+-".join("-" * w for w in widths) + "-+"
240 +
241 + # Print table
242 + print(separator)
243 + print(format_table_row(headers, widths))
244 + print(separator)
245 + for row in rows:
246 + print(format_table_row(row, widths))
247 + print(separator)
248 +
249 +
250 +def format_speedup(value):
251 + """Format speedup value with indicator if it's faster or slower."""
252 + return f"{value:.2f}x {'faster' if value > 1.0 else 'slower'}"
253 +
254 +
255 +def run_benchmarks(verbose: bool = False):
256 + """Run benchmarks for a set of common shapes."""
257 + print("===== STARTING FP8 GEMM BENCHMARK =====")
258 +
259 + # Make sure we're using the GPU
260 + if not torch.cuda.is_available():
261 + print("CUDA not available! Tests require GPU.")
262 + return
263 +
264 + # Print system information
265 + print(f"PyTorch version: {torch.__version__}")
266 + print(f"CUDA version: {torch.version.cuda}")
267 + print(f"Triton version: {triton.__version__}")
268 + print(f"Using device: {torch.cuda.get_device_name()}")
269 +
270 + # Enable TF32 for better performance
271 + torch.backends.cuda.matmul.allow_tf32 = True
272 + torch.backends.cudnn.allow_tf32 = True
273 +
274 + # Set seeds for reproducibility
275 + torch.manual_seed(42)
276 + torch.cuda.manual_seed(42)
277 +
278 + # Define benchmark shapes (m, n, k)
279 + shapes = [
280 + (8, 4096, 7168),
281 + (8, 7168, 18432),
282 + (8, 18432, 7168),
283 + (64, 4096, 7168),
284 + (64, 7168, 18432),
285 + (64, 18432, 7168),
286 + (64, 24576, 1536),
287 + (64, 32768, 512),
288 + (64, 7168, 16384),
289 + (128, 4096, 7168),
290 + (128, 7168, 18432),
291 + (128, 18432, 7168),
292 + (1024, 4096, 7168),
293 + (1024, 18432, 7168),
294 + (2048, 4096, 7168),
295 + (4096, 4096, 7168),
296 + ]
297 + shapes = [
298 + # (64, 2112, 7168),
299 + (64, 24576, 1536),
300 + (64, 32768, 512),
301 + (64, 7168, 16384),
302 + (64, 4096, 7168),
303 + (64, 7168, 2048),
304 + # (128, 2112, 7168),
305 + (128, 24576, 1536),
306 + (128, 32768, 512),
307 + (128, 7168, 16384),
308 + (128, 4096, 7168),
309 + (128, 7168, 2048),
310 + # (4096, 2112, 7168),
311 + (4096, 24576, 1536),
312 + (4096, 32768, 512),
313 + (4096, 7168, 16384),
314 + (4096, 4096, 7168),
315 + (4096, 7168, 2048),
316 + ]
317 +
318 + all_results = []
319 + for m, n, k in shapes:
320 + result = benchmark_shape(m, n, k, verbose=verbose)
321 + all_results.append(result)
322 +
323 + # Print results in a nicely formatted table
324 + print("\n===== PERFORMANCE COMPARISON =====")
325 +
326 + # Print DeepGEMM table
327 + deepgemm_headers = ["m", "n", "k", "Time (μs)", "TFLOPS", "GB/s"]
328 + deepgemm_rows = []
329 + for result in all_results:
330 + shape = result["shape"]
331 + impl_data = result["implementations"]["DeepGEMM"]
332 + deepgemm_rows.append([
333 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
334 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}"
335 + ])
336 +
337 + print_table(deepgemm_headers,
338 + deepgemm_rows,
339 + title="DeepGEMM Implementation:")
340 +
341 + # Print vLLM Triton table
342 + triton_headers = [
343 + "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"
344 + ]
345 + triton_rows = []
346 + for result in all_results:
347 + shape = result["shape"]
348 + impl_data = result["implementations"]["vLLM Triton"]
349 + speedup = impl_data.get("speedup_vs_deepgemm", 1.0)
350 + triton_rows.append([
351 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
352 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
353 + format_speedup(speedup)
354 + ])
355 +
356 + print_table(triton_headers,
357 + triton_rows,
358 + title="vLLM Triton Implementation:")
359 +
360 + # Print vLLM CUTLASS table
361 + cutlass_headers = [
362 + "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM",
363 + "vs Triton"
364 + ]
365 + cutlass_rows = []
366 + for result in all_results:
367 + shape = result["shape"]
368 + impl_data = result["implementations"]["vLLM CUTLASS"]
369 + vs_deepgemm = impl_data.get("speedup_vs_deepgemm", 1.0)
370 + vs_triton = impl_data.get("speedup_vs_triton", 1.0)
371 + cutlass_rows.append([
372 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
373 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
374 + format_speedup(vs_deepgemm),
375 + format_speedup(vs_triton)
376 + ])
377 +
378 + print_table(cutlass_headers,
379 + cutlass_rows,
380 + title="vLLM CUTLASS Implementation:")
381 +
382 + # Calculate and print averages
383 + print("\n===== AVERAGE PERFORMANCE =====")
384 +
385 + implementations = ["DeepGEMM", "vLLM Triton", "vLLM CUTLASS"]
386 + avg_metrics = {
387 + impl: {
388 + "tflops": 0,
389 + "gb_s": 0,
390 + "time_ms": 0
391 + }
392 + for impl in implementations
393 + }
394 +
395 + for result in all_results:
396 + for impl in implementations:
397 + impl_data = result["implementations"][impl]
398 + avg_metrics[impl]["tflops"] += impl_data["tflops"]
399 + avg_metrics[impl]["gb_s"] += impl_data["gb_s"]
400 + avg_metrics[impl]["time_ms"] += impl_data["time_ms"]
401 +
402 + num_shapes = len(all_results)
403 + avg_headers = ["Implementation", "Avg TFLOPS", "Avg GB/s", "Avg Time (ms)"]
404 + avg_rows = []
405 +
406 + for impl in implementations:
407 + avg_tflops = avg_metrics[impl]["tflops"] / num_shapes
408 + avg_mem_bw = avg_metrics[impl]["gb_s"] / num_shapes
409 + avg_time = avg_metrics[impl]["time_ms"] / num_shapes
410 + avg_rows.append([
411 + impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"
412 + ])
413 +
414 + print_table(avg_headers, avg_rows)
415 +
416 + # Calculate average speedups
417 + avg_speedups = {
418 + "DeepGEMM vs vLLM Triton": 0,
419 + "DeepGEMM vs vLLM CUTLASS": 0,
420 + "vLLM CUTLASS vs vLLM Triton": 0
421 + }
422 +
423 + for result in all_results:
424 + deepgemm_time = result["implementations"]["DeepGEMM"]["time_ms"]
425 + vllm_triton_time = result["implementations"]["vLLM Triton"]["time_ms"]
426 + vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"][
427 + "time_ms"]
428 +
429 + avg_speedups[
430 + "DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
431 + avg_speedups[
432 + "DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
433 + avg_speedups[
434 + "vLLM CUTLASS vs vLLM Triton"] += vllm_triton_time / vllm_cutlass_time
435 +
436 + print("\n===== AVERAGE SPEEDUPS =====")
437 + speedup_headers = ["Comparison", "Speedup"]
438 + speedup_rows = []
439 + for comparison, total in avg_speedups.items():
440 + avg_speedup = total / num_shapes
441 + status = "faster" if avg_speedup > 1 else "slower"
442 + speedup_rows.append([comparison, f"{avg_speedup:.2f}x {status}"])
443 +
444 + print_table(speedup_headers, speedup_rows)
445 +
446 + # Average accuracy comparison
447 + print("\n===== ACCURACY COMPARISON =====")
448 + avg_diff = {impl: 0 for impl in implementations}
449 +
450 + for result in all_results:
451 + for impl in implementations:
452 + avg_diff[impl] += result["implementations"][impl]["diff"][
453 + "Reference"]
454 +
455 + diff_headers = ["Implementation", "Avg Diff vs Reference"]
456 + diff_rows = []
457 + for impl in implementations:
458 + diff_rows.append([impl, f"{avg_diff[impl] / num_shapes:.6f}"])
459 +
460 + print_table(diff_headers, diff_rows)
461 +
462 +
463 +if __name__ == "__main__":
464 + run_benchmarks(verbose=False)