| 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_prefill( |
| 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 | max_q_len = max_kv_len = max_seq_len |
| 49 | |
| 50 | num_qo_heads, num_kv_heads = num_heads |
| 51 | assert num_qo_heads % num_kv_heads == 0 |
| 52 | |
| 53 | sm_scale = float(1.0 / (head_size**0.5)) |
| 54 | |
| 55 | # large number to reduce kv_cache reuse |
| 56 | NUM_BLOCKS = int(256000 / block_size) |
| 57 | |
| 58 | kv_cache_shape = None |
| 59 | if kv_layout == "NHD": |
| 60 | kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size) |
| 61 | elif kv_layout == "HND": |
| 62 | kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size) |
| 63 | else: |
| 64 | raise ValueError(f"Invalid kv_layout: {kv_layout}") |
| 65 | |
| 66 | q_lens = torch.randint(1, max_q_len, (batch_size,), dtype=torch.int32) |
| 67 | q_lens[-1] = max_q_len |
| 68 | q_indptr = torch.cat( |
| 69 | [ |
| 70 | torch.tensor([0], dtype=torch.int32), |
| 71 | torch.cumsum(q_lens, dim=0, dtype=torch.int32), |
| 72 | ] |
| 73 | ) |
| 74 | |
| 75 | # Always using 1.0 scale to reflect the real perf in benchmarking |
| 76 | q_scale = 1.0 |
| 77 | ref_query = torch.randn( |
| 78 | torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype |
| 79 | ) |
| 80 | if q_quant_dtype == FP8_DTYPE: |
| 81 | query, _ = to_float8(ref_query) |
| 82 | else: |
| 83 | query = ref_query |
| 84 | |
| 85 | kv_lens = torch.randint(0, max_kv_len, (batch_size,), dtype=torch.int32) |
| 86 | kv_lens[-1] = max_kv_len |
| 87 | |
| 88 | seq_lens = kv_lens + q_lens |
| 89 | max_seq_len = torch.max(seq_lens).item() |
| 90 | |
| 91 | # Always using 1.0 scale to reflect the real perf in benchmarking |
| 92 | k_scale = v_scale = 1.0 |
| 93 | ref_kv_cache = torch.randn(kv_cache_shape, dtype=dtype) |
| 94 | if kv_quant_dtype == FP8_DTYPE: |
| 95 | kv_cache, _ = to_float8(ref_kv_cache) |
| 96 | else: |
| 97 | kv_cache = ref_kv_cache |
| 98 | |
| 99 | max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size |
| 100 | block_tables = torch.randint( |
| 101 | 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32 |
| 102 | ) |
| 103 | kv_indptr = [0] |
| 104 | kv_indices = [] |
| 105 | kv_last_page_lens = [] |
| 106 | for i in range(batch_size): |
| 107 | seq_len = seq_lens[i] |
| 108 | assert seq_len > 0 |
| 109 | num_blocks = (seq_len + block_size - 1) // block_size |
| 110 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 111 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 112 | kv_last_page_len = seq_len % block_size |
| 113 | if kv_last_page_len == 0: |
| 114 | kv_last_page_len = block_size |
| 115 | kv_last_page_lens.append(kv_last_page_len) |
| 116 | |
| 117 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 118 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 119 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 120 | workspace_buffer = torch.zeros(1024 * 1024 * 1024, dtype=torch.int8) |
| 121 | |
| 122 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( |
| 123 | workspace_buffer, kv_layout |
| 124 | ) |
| 125 | wrapper.plan( |
| 126 | q_indptr, |
| 127 | kv_indptr, |
| 128 | kv_indices, |
| 129 | kv_last_page_lens, |
| 130 | num_qo_heads, |
| 131 | num_kv_heads, |
| 132 | head_size, |
| 133 | block_size, |
| 134 | causal=True, |
| 135 | sm_scale=sm_scale, |
| 136 | q_data_type=dtype, |
| 137 | kv_data_type=dtype, |
| 138 | ) |
| 139 | |
| 140 | def time_fn(fn, warmup=10, trials=20): |
| 141 | torch.accelerator.synchronize() |
| 142 | start = torch.Event(enable_timing=True) |
| 143 | end = torch.Event(enable_timing=True) |
| 144 | times = [] |
| 145 | for i in range(warmup): |
| 146 | fn() |
| 147 | for i in range(trials): |
| 148 | start.record() |
| 149 | fn() |
| 150 | end.record() |
| 151 | torch.accelerator.synchronize() |
| 152 | times.append(start.elapsed_time(end)) # ms |
| 153 | return sum(times) / len(times), torch.std(torch.tensor(times)) |
| 154 | |
| 155 | o_scale = 1.0 |
| 156 | o_sf_scale = None |
| 157 | output_baseline = torch.empty(ref_query.shape, dtype=dtype) |
| 158 | if o_quant_dtype == FP4_DTYPE: |
| 159 | o_sf_scale = 500.0 |
| 160 | output_trtllm = flashinfer.utils.FP4Tensor( |
| 161 | torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8), |
| 162 | torch.empty( |
| 163 | ( |
| 164 | round_up(query.shape[0], 128), |
| 165 | round_up(query.shape[1] * query.shape[2] // 16, 4), |
| 166 | ), |
| 167 | dtype=torch.float8_e4m3fn, |
| 168 | ), |
| 169 | ) |
| 170 | else: |
| 171 | output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype) |
| 172 | |
| 173 | def baseline_prefill(): |
| 174 | return wrapper.run( |
| 175 | ref_query, |
| 176 | ref_kv_cache, |
| 177 | k_scale=k_scale, |
| 178 | v_scale=v_scale, |
| 179 | out=output_baseline, |
| 180 | ) |
| 181 | |
| 182 | def trtllm_prefill(): |
| 183 | return flashinfer.prefill.trtllm_batch_context_with_kv_cache( |
| 184 | query=query, |
| 185 | kv_cache=kv_cache, |
| 186 | workspace_buffer=workspace_buffer, |
| 187 | block_tables=block_tables, |
| 188 | seq_lens=seq_lens, |
| 189 | max_q_len=max_q_len, |
| 190 | max_kv_len=max_seq_len, |
| 191 | bmm1_scale=q_scale * k_scale * sm_scale, |
| 192 | bmm2_scale=v_scale / o_scale, |
| 193 | batch_size=batch_size, |
| 194 | cum_seq_lens_q=q_indptr, |
| 195 | cum_seq_lens_kv=kv_indptr, |
| 196 | o_sf_scale=o_sf_scale, |
| 197 | out=output_trtllm, |
| 198 | ) |
| 199 | |
| 200 | baseline_mean, baseline_std = time_fn(baseline_prefill) |
| 201 | trtllm_mean, trtllm_std = time_fn(trtllm_prefill) |
| 202 | |
| 203 | # Calculate percentage speedup (positive means TRT is faster) |
| 204 | speedup_percent = (baseline_mean - trtllm_mean) / baseline_mean |
| 205 | |
| 206 | print( |
| 207 | f"\t{batch_size}\t{max_seq_len}\t{trtllm_mean:8.3f}\t{trtllm_std.item():8.3f}" |
| 208 | f"\t{baseline_mean:8.3f}\t{baseline_std.item():8.3f}\t{speedup_percent:8.3f}" |
| 209 | ) |
| 210 | |
| 211 | # Return results for CSV writing |
| 212 | return { |
| 213 | "batch_size": batch_size, |
| 214 | "trtllm_mean": trtllm_mean, |
| 215 | "trtllm_std": trtllm_std.item(), |
| 216 | "baseline_mean": baseline_mean, |
| 217 | "baseline_std": baseline_std.item(), |
| 218 | "speedup_percent": speedup_percent, |
| 219 | "q_dtype": str(q_quant_dtype), |
| 220 | "kv_cache_dtype": str(kv_quant_dtype), |
| 221 | "output_dtype": str(o_quant_dtype), |
| 222 | "block_size": block_size, |
| 223 | "num_kv_heads": num_kv_heads, |
| 224 | "head_size": head_size, |
| 225 | "max_seq_len": max_seq_len, |
| 226 | } |
| 227 | |
| 228 | |
| 229 | def write_results_to_csv(results, filename=None): |
| 230 | """Write benchmark results to CSV file.""" |
| 231 | if filename is None: |
| 232 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 233 | filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv" |
| 234 | |
| 235 | fieldnames = [ |
| 236 | "batch_size", |
| 237 | "trtllm_mean", |
| 238 | "trtllm_std", |
| 239 | "baseline_mean", |
| 240 | "baseline_std", |
| 241 | "speedup_percent", |
| 242 | "q_dtype", |
| 243 | "kv_cache_dtype", |
| 244 | "output_dtype", |
| 245 | "block_size", |
| 246 | "num_kv_heads", |
| 247 | "head_size", |
| 248 | "max_seq_len", |
| 249 | ] |
| 250 | |
| 251 | file_exists = os.path.exists(filename) |
| 252 | |
| 253 | with open(filename, "a", newline="") as csvfile: |
| 254 | writer = csv.DictWriter(csvfile, fieldnames=fieldnames) |
| 255 | |
| 256 | if not file_exists: |
| 257 | writer.writeheader() |
| 258 | |
| 259 | for result in results: |
| 260 | writer.writerow(result) |
| 261 | |
| 262 | print(f"Results written to {filename}") |
| 263 | |
| 264 | |
| 265 | if __name__ == "__main__": |
| 266 | batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256] |
| 267 | max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072] |
| 268 | all_results = [] |
| 269 | |
| 270 | dtype = torch.bfloat16 |
| 271 | quant_dtypes = [ |
| 272 | # (q_quant_dtype, kv_quant_dtype, o_quant_dtype) |
| 273 | (None, None, None), |
| 274 | (FP8_DTYPE, FP8_DTYPE, None), |
| 275 | (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE), |
| 276 | (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE), |
| 277 | ] |
| 278 | |
| 279 | for quant_dtype in quant_dtypes: |
| 280 | q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtype |
| 281 | q_quant_dtype = q_quant_dtype or dtype |
| 282 | kv_quant_dtype = kv_quant_dtype or dtype |
| 283 | o_quant_dtype = o_quant_dtype or dtype |
| 284 | |
| 285 | print( |
| 286 | f"Running benchmark for q_dtype = {q_quant_dtype}, " |
| 287 | f"kv_cache_dtype: {kv_quant_dtype}, " |
| 288 | f"output_dtype: {o_quant_dtype}" |
| 289 | ) |
| 290 | print( |
| 291 | "\tbatch_size\tmax_seq_len\ttrtllm_mean\ttrtllm_std\tbaseline_mean\t" |
| 292 | "baseline_std\tspeedup_percent" |
| 293 | ) |
| 294 | for max_seq_len in max_seq_lens: |
| 295 | for bs in batch_sizes: |
| 296 | result = benchmark_prefill( |
| 297 | dtype=dtype, |
| 298 | quant_dtypes=quant_dtype, |
| 299 | batch_size=bs, |
| 300 | max_seq_len=max_seq_len, |
| 301 | ) |
| 302 | all_results.append(result) |
| 303 | |
| 304 | # Write all results to CSV |
| 305 | write_results_to_csv(all_results) |