| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import csv |
| 5 | import os |
| 6 | from datetime import datetime |
| 7 | |
| 8 | import flashinfer |
| 9 | import torch |
| 10 | |
| 11 | from vllm.utils.math_utils import round_up |
| 12 | |
| 13 | FLOAT32_BYTES = torch.finfo(torch.float).bits // 8 |
| 14 | FP8_DTYPE = torch.float8_e4m3fn |
| 15 | FP4_DTYPE = torch.uint8 |
| 16 | |
| 17 | |
| 18 | def to_float8(x, dtype=torch.float8_e4m3fn): |
| 19 | finfo = torch.finfo(dtype) |
| 20 | min_val, max_val = x.aminmax() |
| 21 | amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) |
| 22 | scale = finfo.max / amax * 0.1 |
| 23 | x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) |
| 24 | return x_scl_sat.to(dtype), scale.float().reciprocal() |
| 25 | |
| 26 | |
| 27 | @torch.no_grad() |
| 28 | def benchmark_decode( |
| 29 | dtype: torch.dtype, |
| 30 | quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None], |
| 31 | batch_size: int, |
| 32 | max_seq_len: int, |
| 33 | num_heads: tuple[int, int] = (64, 8), |
| 34 | head_size: int = 128, |
| 35 | kv_layout: str = "HND", |
| 36 | block_size: int = 16, |
| 37 | warmup: int = 10, |
| 38 | trials: int = 20, |
| 39 | ): |
| 40 | torch.set_default_device("cuda") |
| 41 | torch.manual_seed(0) |
| 42 | |
| 43 | q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes |
| 44 | q_quant_dtype = q_quant_dtype or dtype |
| 45 | kv_quant_dtype = kv_quant_dtype or dtype |
| 46 | o_quant_dtype = o_quant_dtype or dtype |
| 47 | |
| 48 | num_qo_heads, num_kv_heads = num_heads |
| 49 | assert num_qo_heads % num_kv_heads == 0 |
| 50 | |
| 51 | sm_scale = float(1.0 / (head_size**0.5)) |
| 52 | |
| 53 | # large number to reduce kv_cache reuse |
| 54 | NUM_BLOCKS = int(256000 / block_size) |
| 55 | |
| 56 | kv_cache_shape = None |
| 57 | if kv_layout == "NHD": |
| 58 | kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size) |
| 59 | elif kv_layout == "HND": |
| 60 | kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size) |
| 61 | else: |
| 62 | raise ValueError(f"Invalid kv_layout: {kv_layout}") |
| 63 | |
| 64 | # Always using 1.0 scale to reflect the real perf in benchmarking |
| 65 | q_scale = 1.0 |
| 66 | ref_query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype) |
| 67 | if q_quant_dtype == FP8_DTYPE: |
| 68 | query, _ = to_float8(ref_query) |
| 69 | else: |
| 70 | query = ref_query |
| 71 | |
| 72 | kv_lens = torch.randint(1, max_seq_len, (batch_size,), dtype=torch.int32) |
| 73 | kv_lens[-1] = max_seq_len |
| 74 | |
| 75 | seq_lens = kv_lens |
| 76 | max_seq_len = torch.max(seq_lens).item() |
| 77 | |
| 78 | # Always using 1.0 scale to reflect the real perf in benchmarking |
| 79 | k_scale = v_scale = 1.0 |
| 80 | ref_kv_cache = torch.randn(kv_cache_shape, dtype=dtype) |
| 81 | if kv_quant_dtype == FP8_DTYPE: |
| 82 | kv_cache, _ = to_float8(ref_kv_cache) |
| 83 | else: |
| 84 | kv_cache = ref_kv_cache |
| 85 | |
| 86 | max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size |
| 87 | block_tables = torch.randint( |
| 88 | 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32 |
| 89 | ) |
| 90 | kv_indptr = [0] |
| 91 | kv_indices = [] |
| 92 | kv_last_page_lens = [] |
| 93 | for i in range(batch_size): |
| 94 | seq_len = seq_lens[i] |
| 95 | assert seq_len > 0 |
| 96 | num_blocks = (seq_len + block_size - 1) // block_size |
| 97 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 98 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 99 | kv_last_page_len = seq_len % block_size |
| 100 | if kv_last_page_len == 0: |
| 101 | kv_last_page_len = block_size |
| 102 | kv_last_page_lens.append(kv_last_page_len) |
| 103 | |
| 104 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 105 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 106 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 107 | workspace_buffer = torch.zeros(1024 * 1024 * 1024, dtype=torch.int8) |
| 108 | |
| 109 | wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( |
| 110 | workspace_buffer, |
| 111 | kv_layout, |
| 112 | use_tensor_cores=True, |
| 113 | ) |
| 114 | wrapper.plan( |
| 115 | kv_indptr, |
| 116 | kv_indices, |
| 117 | kv_last_page_lens, |
| 118 | num_qo_heads, |
| 119 | num_kv_heads, |
| 120 | head_size, |
| 121 | block_size, |
| 122 | "NONE", |
| 123 | sm_scale=sm_scale, |
| 124 | q_data_type=dtype, |
| 125 | kv_data_type=dtype, |
| 126 | ) |
| 127 | |
| 128 | def time_fn(fn, warmup=10, trials=20): |
| 129 | torch.accelerator.synchronize() |
| 130 | start = torch.Event(enable_timing=True) |
| 131 | end = torch.Event(enable_timing=True) |
| 132 | times = [] |
| 133 | for i in range(warmup): |
| 134 | fn() |
| 135 | for i in range(trials): |
| 136 | start.record() |
| 137 | fn() |
| 138 | end.record() |
| 139 | torch.accelerator.synchronize() |
| 140 | times.append(start.elapsed_time(end)) # ms |
| 141 | return sum(times) / len(times), torch.std(torch.tensor(times)) |
| 142 | |
| 143 | o_scale = 1.0 |
| 144 | o_sf_scale = None |
| 145 | output_baseline = torch.empty(ref_query.shape, dtype=dtype) |
| 146 | if o_quant_dtype == FP4_DTYPE: |
| 147 | o_sf_scale = 500.0 |
| 148 | output_trtllm = flashinfer.utils.FP4Tensor( |
| 149 | torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8), |
| 150 | torch.empty( |
| 151 | ( |
| 152 | round_up(query.shape[0], 128), |
| 153 | round_up(query.shape[1] * query.shape[2] // 16, 4), |
| 154 | ), |
| 155 | dtype=torch.float8_e4m3fn, |
| 156 | ), |
| 157 | ) |
| 158 | else: |
| 159 | output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype) |
| 160 | |
| 161 | def baseline_decode(): |
| 162 | return wrapper.run( |
| 163 | ref_query, |
| 164 | ref_kv_cache, |
| 165 | k_scale=k_scale, |
| 166 | v_scale=v_scale, |
| 167 | out=output_baseline, |
| 168 | ) |
| 169 | |
| 170 | def trtllm_decode(): |
| 171 | return flashinfer.decode.trtllm_batch_decode_with_kv_cache( |
| 172 | query=query, |
| 173 | kv_cache=kv_cache, |
| 174 | workspace_buffer=workspace_buffer, |
| 175 | block_tables=block_tables, |
| 176 | seq_lens=seq_lens, |
| 177 | max_seq_len=max_seq_len, |
| 178 | bmm1_scale=q_scale * k_scale * sm_scale, |
| 179 | bmm2_scale=v_scale / o_scale, |
| 180 | o_sf_scale=o_sf_scale, |
| 181 | out=output_trtllm, |
| 182 | ) |
| 183 | |
| 184 | baseline_mean, baseline_std = time_fn(baseline_decode) |
| 185 | trtllm_mean, trtllm_std = time_fn(trtllm_decode) |
| 186 | |
| 187 | # Calculate percentage speedup (positive means TRT is faster) |
| 188 | speedup_percent = (baseline_mean - trtllm_mean) / baseline_mean |
| 189 | |
| 190 | print( |
| 191 | f"\t{batch_size}\t{max_seq_len}\t{trtllm_mean:.3f}\t{trtllm_std.item():.3f}" |
| 192 | f"\t{baseline_mean:.3f}\t{baseline_std.item():.3f}\t{speedup_percent:.3f}" |
| 193 | ) |
| 194 | |
| 195 | # Return results for CSV writing |
| 196 | return { |
| 197 | "batch_size": batch_size, |
| 198 | "trtllm_mean": trtllm_mean, |
| 199 | "trtllm_std": trtllm_std.item(), |
| 200 | "baseline_mean": baseline_mean, |
| 201 | "baseline_std": baseline_std.item(), |
| 202 | "speedup_percent": speedup_percent, |
| 203 | "q_dtype": str(q_quant_dtype), |
| 204 | "kv_cache_dtype": str(kv_quant_dtype), |
| 205 | "output_dtype": str(o_quant_dtype), |
| 206 | "block_size": block_size, |
| 207 | "num_kv_heads": num_kv_heads, |
| 208 | "head_size": head_size, |
| 209 | "max_seq_len": max_seq_len, |
| 210 | } |
| 211 | |
| 212 | |
| 213 | def write_results_to_csv(results, filename=None): |
| 214 | """Write benchmark results to CSV file.""" |
| 215 | if filename is None: |
| 216 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 217 | filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv" |
| 218 | |
| 219 | fieldnames = [ |
| 220 | "batch_size", |
| 221 | "trtllm_mean", |
| 222 | "trtllm_std", |
| 223 | "baseline_mean", |
| 224 | "baseline_std", |
| 225 | "speedup_percent", |
| 226 | "q_dtype", |
| 227 | "kv_cache_dtype", |
| 228 | "output_dtype", |
| 229 | "block_size", |
| 230 | "num_kv_heads", |
| 231 | "head_size", |
| 232 | "max_seq_len", |
| 233 | ] |
| 234 | |
| 235 | file_exists = os.path.exists(filename) |
| 236 | |
| 237 | with open(filename, "a", newline="") as csvfile: |
| 238 | writer = csv.DictWriter(csvfile, fieldnames=fieldnames) |
| 239 | |
| 240 | if not file_exists: |
| 241 | writer.writeheader() |
| 242 | |
| 243 | for result in results: |
| 244 | writer.writerow(result) |
| 245 | |
| 246 | print(f"Results written to {filename}") |
| 247 | |
| 248 | |
| 249 | if __name__ == "__main__": |
| 250 | batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256] |
| 251 | max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072] |
| 252 | all_results = [] |
| 253 | |
| 254 | dtype = torch.bfloat16 |
| 255 | quant_dtypes = [ |
| 256 | # (q_quant_dtype, kv_quant_dtype, o_quant_dtype) |
| 257 | (None, None, None), |
| 258 | (None, FP8_DTYPE, None), |
| 259 | (FP8_DTYPE, FP8_DTYPE, None), |
| 260 | (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE), |
| 261 | (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE), |
| 262 | ] |
| 263 | |
| 264 | for quant_dtype in quant_dtypes: |
| 265 | q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtype |
| 266 | q_quant_dtype = q_quant_dtype or dtype |
| 267 | kv_quant_dtype = kv_quant_dtype or dtype |
| 268 | o_quant_dtype = o_quant_dtype or dtype |
| 269 | |
| 270 | print( |
| 271 | f"Running benchmark for q_dtype = {q_quant_dtype}, " |
| 272 | f"kv_cache_dtype: {kv_quant_dtype}, " |
| 273 | f"output_dtype: {o_quant_dtype}" |
| 274 | ) |
| 275 | print( |
| 276 | "\tbatch_size\tmax_seq_len\ttrtllm_mean\ttrtllm_std\tbaseline_mean\t" |
| 277 | "baseline_std\tspeedup_percent" |
| 278 | ) |
| 279 | for max_seq_len in max_seq_lens: |
| 280 | for bs in batch_sizes: |
| 281 | result = benchmark_decode( |
| 282 | dtype=dtype, |
| 283 | quant_dtypes=quant_dtype, |
| 284 | batch_size=bs, |
| 285 | max_seq_len=max_seq_len, |
| 286 | ) |
| 287 | all_results.append(result) |
| 288 | |
| 289 | # Write all results to CSV |
| 290 | write_results_to_csv(all_results) |