| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import copy |
| 6 | import itertools |
| 7 | import math |
| 8 | import os |
| 9 | import pickle as pkl |
| 10 | import time |
| 11 | from collections.abc import Callable, Iterable |
| 12 | from dataclasses import dataclass |
| 13 | from itertools import product |
| 14 | |
| 15 | import pandas as pd |
| 16 | import torch |
| 17 | import torch.utils.benchmark as TBenchmark |
| 18 | from torch.utils.benchmark import Measurement as TMeasurement |
| 19 | from weight_shapes import WEIGHT_SHAPES |
| 20 | |
| 21 | from vllm import _custom_ops as ops |
| 22 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 23 | GPTQ_MARLIN_MAX_PARALLEL, |
| 24 | GPTQ_MARLIN_MIN_THREAD_N, |
| 25 | marlin_permute_scales, |
| 26 | marlin_zero_points, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( |
| 29 | MarlinWorkspace, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 32 | pack_rows, |
| 33 | quantize_weights, |
| 34 | ) |
| 35 | from vllm.scalar_type import ScalarType, scalar_types |
| 36 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 37 | |
| 38 | DEFAULT_MODELS = ["meta-llama/Llama-3-8b", "meta-llama/Llama-2-70b-hf"] |
| 39 | DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024] |
| 40 | DEFAULT_TP_SIZES = [1] |
| 41 | |
| 42 | NVTX_PROFILE = os.environ.get("NVTX_PROFILE", False) |
| 43 | |
| 44 | if NVTX_PROFILE: |
| 45 | import nvtx |
| 46 | |
| 47 | |
| 48 | def terse_type_name(dt): |
| 49 | return { |
| 50 | torch.bfloat16: "bf16", |
| 51 | torch.float16: "fp16", |
| 52 | torch.int8: "int8", |
| 53 | torch.float8_e4m3fn: "fp8", |
| 54 | torch.float: "float", |
| 55 | torch.int: "int", |
| 56 | }[dt] |
| 57 | |
| 58 | |
| 59 | @dataclass |
| 60 | class BenchmarkTensors: |
| 61 | w_ref: torch.Tensor |
| 62 | a: torch.Tensor |
| 63 | |
| 64 | w_q: torch.Tensor |
| 65 | group_size: int | None |
| 66 | wtype: ScalarType |
| 67 | w_g_s: torch.Tensor |
| 68 | w_g_zp: torch.Tensor | None |
| 69 | w_ch_s: torch.Tensor | None |
| 70 | w_tok_s: torch.Tensor | None |
| 71 | |
| 72 | |
| 73 | @dataclass |
| 74 | class TypeConfig: |
| 75 | act_type: torch.dtype |
| 76 | weight_type: ScalarType |
| 77 | output_type: torch.dtype | None |
| 78 | group_scale_type: torch.dtype | None |
| 79 | group_zero_type: torch.dtype | None |
| 80 | channel_scale_type: torch.dtype | None |
| 81 | token_scale_type: torch.dtype | None |
| 82 | |
| 83 | |
| 84 | def rand_data(shape, dtype=torch.float16, scale=1): |
| 85 | if dtype.is_floating_point: |
| 86 | return (scale * torch.rand(shape, device="cuda") - 0.3).to(dtype) |
| 87 | else: |
| 88 | return torch.randint(-15, 15, shape, dtype=dtype, device="cuda") |
| 89 | |
| 90 | |
| 91 | def quantize_and_pack( |
| 92 | atype: torch.dtype, |
| 93 | w: torch.Tensor, |
| 94 | wtype: ScalarType, |
| 95 | stype: torch.dtype | None, |
| 96 | group_size: int | None, |
| 97 | zero_points: bool = False, |
| 98 | ): |
| 99 | assert wtype.is_integer(), "TODO: support floating point weights" |
| 100 | |
| 101 | w_ref, w_q, w_s, w_zp = quantize_weights( |
| 102 | w, |
| 103 | wtype, |
| 104 | group_size=group_size, |
| 105 | zero_points=zero_points, |
| 106 | # to match how the kernel applies zps |
| 107 | ref_zero_points_after_scales=True, |
| 108 | ) |
| 109 | |
| 110 | w_q = pack_rows(w_q, wtype.size_bits, *w_q.shape) |
| 111 | return w_ref, w_q, w_s, w_zp |
| 112 | |
| 113 | |
| 114 | def create_bench_tensors( |
| 115 | shape: tuple[int, int, int], types: TypeConfig, group_size: int | None |
| 116 | ) -> list[BenchmarkTensors]: |
| 117 | m, n, k = shape |
| 118 | |
| 119 | # we want to make sure that weights don't fit into L2 cache between runs so |
| 120 | # we construct enough weights to exceed L2 cache, which is 50mb on a H100 |
| 121 | # so we target total weight size > 2*50mb |
| 122 | num_weights = math.ceil( |
| 123 | 2 * 50 * 1024**2 * 8 / (k * n * types.weight_type.size_bits) |
| 124 | ) |
| 125 | |
| 126 | a = rand_data((m, k), types.act_type, scale=5) |
| 127 | |
| 128 | benchmark_tensors: list[BenchmarkTensors] = [] |
| 129 | for _ in range(num_weights): |
| 130 | w = rand_data((k, n), types.act_type, scale=5) |
| 131 | |
| 132 | if types.group_scale_type is not None: |
| 133 | w = w.to(types.group_scale_type) |
| 134 | if w.dtype.itemsize == 1: |
| 135 | w = w.to(torch.float16) |
| 136 | |
| 137 | w_ref, w_q_packed, w_s, w_zp = quantize_and_pack( |
| 138 | a.dtype, |
| 139 | w, |
| 140 | types.weight_type, |
| 141 | types.group_scale_type, |
| 142 | group_size, |
| 143 | types.group_zero_type is not None, |
| 144 | ) |
| 145 | |
| 146 | if not a.dtype.is_floating_point: |
| 147 | aiinfo = torch.iinfo(a.dtype) |
| 148 | w_ref = w_ref.round().clamp(aiinfo.min, aiinfo.max) |
| 149 | |
| 150 | w_ref = w_ref.to(torch.float32) |
| 151 | |
| 152 | w_ch_s = ( |
| 153 | None |
| 154 | if types.channel_scale_type is None |
| 155 | else rand_data((n,), types.channel_scale_type) |
| 156 | ) |
| 157 | w_tok_s = ( |
| 158 | None |
| 159 | if types.token_scale_type is None |
| 160 | else rand_data((m,), types.token_scale_type) |
| 161 | ) |
| 162 | |
| 163 | benchmark_tensors.append( |
| 164 | BenchmarkTensors( |
| 165 | w_ref=w_ref, |
| 166 | a=a, |
| 167 | w_q=w_q_packed, |
| 168 | wtype=types.weight_type, |
| 169 | w_g_s=w_s, |
| 170 | w_g_zp=w_zp, |
| 171 | group_size=group_size, |
| 172 | w_ch_s=w_ch_s, |
| 173 | w_tok_s=w_tok_s, |
| 174 | ) |
| 175 | ) |
| 176 | |
| 177 | return benchmark_tensors |
| 178 | |
| 179 | |
| 180 | def torch_matmul_f16_create_bench_fn(bt: BenchmarkTensors) -> Callable: |
| 181 | a = bt.a |
| 182 | w = bt.w_ref.to(bt.a.dtype) # use float reference tensor |
| 183 | if a.dtype not in [torch.float16, torch.bfloat16]: |
| 184 | a = a.to(torch.float16) |
| 185 | w = w.to(torch.float16) |
| 186 | return lambda: torch.matmul(a, w) |
| 187 | |
| 188 | |
| 189 | def cutlass_scaled_mm_create_bench_fn(bt: BenchmarkTensors) -> Callable: |
| 190 | if bt.w_ch_s is not None and bt.w_tok_s is not None: |
| 191 | scale_a = bt.w_tok_s.to(torch.float32) |
| 192 | scale_b = bt.w_ch_s.to(torch.float32) |
| 193 | else: |
| 194 | scale_a = torch.tensor(1.0, dtype=torch.float32, device=bt.a.device) |
| 195 | scale_b = torch.tensor(1.0, dtype=torch.float32, device=bt.a.device) |
| 196 | w_col_major = bt.w_ref.to(bt.a.dtype).t().contiguous().t() |
| 197 | return lambda: ops.cutlass_scaled_mm( |
| 198 | bt.a, w_col_major, scale_a, scale_b, out_dtype=torch.float16 |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable: |
| 203 | device = bt.a.device |
| 204 | |
| 205 | workspace = MarlinWorkspace( |
| 206 | bt.w_ref.shape[1], GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL |
| 207 | ) |
| 208 | |
| 209 | if bt.w_g_zp is None: |
| 210 | w_zp = torch.empty(0, dtype=torch.int, device=device) |
| 211 | else: |
| 212 | w_zp = marlin_zero_points( |
| 213 | bt.w_g_zp, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits |
| 214 | ) |
| 215 | |
| 216 | if bt.group_size is None: |
| 217 | w_s = torch.tensor([], device="cuda", dtype=torch.half) |
| 218 | else: |
| 219 | w_s = marlin_permute_scales( |
| 220 | bt.w_g_s, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.group_size |
| 221 | ) |
| 222 | |
| 223 | w_q = ops.gptq_marlin_repack( |
| 224 | bt.w_q, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits |
| 225 | ) |
| 226 | |
| 227 | if bt.a.dtype.is_floating_point: |
| 228 | assert bt.w_ch_s is None |
| 229 | assert bt.w_tok_s is None |
| 230 | assert bt.group_size is not None |
| 231 | |
| 232 | fn = lambda: ops.marlin_gemm( |
| 233 | a=bt.a, |
| 234 | c=None, |
| 235 | b_q_weight=w_q, |
| 236 | b_bias=None, |
| 237 | b_scales=w_s, |
| 238 | a_scales=None, |
| 239 | global_scale=None, |
| 240 | b_zeros=w_zp, |
| 241 | workspace=workspace.scratch, |
| 242 | b_q_type=bt.wtype, |
| 243 | size_m=bt.a.shape[0], |
| 244 | size_n=bt.w_ref.shape[1], |
| 245 | size_k=bt.w_ref.shape[0], |
| 246 | is_zp_float=False, |
| 247 | ) |
| 248 | else: |
| 249 | assert bt.a.dtype == torch.int8 |
| 250 | assert bt.wtype == scalar_types.uint4b8 |
| 251 | raise NotImplementedError("QQQ is not supported anymore") |
| 252 | |
| 253 | return fn |
| 254 | |
| 255 | |
| 256 | def machete_create_bench_fn( |
| 257 | bt: BenchmarkTensors, out_type=torch.dtype, schedule=None |
| 258 | ) -> Callable: |
| 259 | w_q = bt.w_q.t().contiguous().t() # make col major |
| 260 | w_q = ops.machete_prepack_B( |
| 261 | w_q, bt.a.dtype, bt.wtype, None if bt.w_g_s is None else bt.w_g_s.dtype |
| 262 | ) |
| 263 | |
| 264 | w_g_zp = bt.w_g_zp |
| 265 | if w_g_zp is not None: |
| 266 | w_g_zp = -1 * bt.w_g_s * (w_g_zp.to(bt.w_g_s.dtype)) |
| 267 | |
| 268 | return lambda: ops.machete_mm( |
| 269 | a=bt.a, |
| 270 | b_q=w_q, |
| 271 | b_type=bt.wtype, |
| 272 | b_group_scales=bt.w_g_s, |
| 273 | b_group_zeros=w_g_zp, |
| 274 | b_group_size=bt.group_size, |
| 275 | b_channel_scales=bt.w_ch_s, |
| 276 | a_token_scales=bt.w_tok_s, |
| 277 | out_type=out_type, |
| 278 | schedule=schedule, |
| 279 | ) |
| 280 | |
| 281 | |
| 282 | def cutlass_w4a8_create_bench_fn( |
| 283 | bt: BenchmarkTensors, out_type=torch.dtype, schedule=None |
| 284 | ) -> Callable: |
| 285 | w_q = bt.w_q.t().contiguous().t() # make col major |
| 286 | w_q = ops.cutlass_encode_and_reorder_int4b(w_q) |
| 287 | # expects fp8 scales |
| 288 | w_s = ops.cutlass_pack_scale_fp8(bt.w_g_s.to(torch.float8_e4m3fn)) |
| 289 | |
| 290 | return lambda: ops.cutlass_w4a8_mm( |
| 291 | a=bt.a, |
| 292 | b_q=w_q, |
| 293 | b_group_scales=w_s, |
| 294 | b_group_size=bt.group_size, |
| 295 | b_channel_scales=bt.w_ch_s, |
| 296 | a_token_scales=bt.w_tok_s, |
| 297 | maybe_schedule=schedule, |
| 298 | ) |
| 299 | |
| 300 | |
| 301 | # impl |
| 302 | |
| 303 | # bench |
| 304 | |
| 305 | |
| 306 | def bench_fns(label: str, sub_label: str, description: str, fns: list[Callable]): |
| 307 | min_run_time = 1 if not NVTX_PROFILE else 0.1 |
| 308 | res = TBenchmark.Timer( |
| 309 | stmt=""" |
| 310 | for fn in fns: |
| 311 | fn() |
| 312 | """, |
| 313 | globals={"fns": fns}, |
| 314 | label=label, |
| 315 | sub_label=sub_label, |
| 316 | description=description, |
| 317 | ).blocked_autorange(min_run_time=min_run_time) |
| 318 | |
| 319 | if NVTX_PROFILE: |
| 320 | with ( |
| 321 | nvtx.annotate("mm-bench"), |
| 322 | nvtx.annotate(f"{label}|{sub_label}|{description}"), |
| 323 | ): |
| 324 | fns[0]() |
| 325 | |
| 326 | return res |
| 327 | |
| 328 | |
| 329 | _SWEEP_SCHEDULES_RESULTS: pd.DataFrame | None = None |
| 330 | _SWEEP_SCHEDULES_RESULTS_CSV: str | None = None |
| 331 | |
| 332 | |
| 333 | def bench( |
| 334 | types: TypeConfig, |
| 335 | group_size: int, |
| 336 | m: int, |
| 337 | k: int, |
| 338 | n: int, |
| 339 | label: str, |
| 340 | sub_label: str, |
| 341 | sweep_schedules: bool = True, |
| 342 | ) -> list[TMeasurement]: |
| 343 | benchmark_tensors = create_bench_tensors((m, n, k), types, group_size) |
| 344 | sub_label += f", L={len(benchmark_tensors)}" |
| 345 | |
| 346 | name_type_string = f"W{types.weight_type}" + f"-A{terse_type_name(types.act_type)}" |
| 347 | if types.group_scale_type is not None: |
| 348 | name_type_string += f"-GS{terse_type_name(types.group_scale_type)}" |
| 349 | if types.group_zero_type is not None: |
| 350 | name_type_string += f"-GZ{terse_type_name(types.group_zero_type)}" |
| 351 | if group_size is not None: |
| 352 | name_type_string += f"-G{group_size}" |
| 353 | if types.channel_scale_type is not None: |
| 354 | name_type_string += f"-CS{terse_type_name(types.channel_scale_type)}" |
| 355 | if types.token_scale_type is not None: |
| 356 | name_type_string += f"-TS{terse_type_name(types.token_scale_type)}" |
| 357 | |
| 358 | timers = [] |
| 359 | # pytorch impl |
| 360 | timers.append( |
| 361 | bench_fns( |
| 362 | label, |
| 363 | sub_label, |
| 364 | "torch.matmul (fp16)", |
| 365 | [torch_matmul_f16_create_bench_fn(bt) for bt in benchmark_tensors], |
| 366 | ) |
| 367 | ) |
| 368 | |
| 369 | if types.act_type == torch.int8 or types.act_type == torch.float8_e4m3fn: |
| 370 | timers.append( |
| 371 | bench_fns( |
| 372 | label, |
| 373 | sub_label, |
| 374 | f"cutlass_scaled_mm ({terse_type_name(types.act_type)})", |
| 375 | [cutlass_scaled_mm_create_bench_fn(bt) for bt in benchmark_tensors], |
| 376 | ) |
| 377 | ) |
| 378 | |
| 379 | if types.act_type != torch.float8_e4m3fn: |
| 380 | timers.append( |
| 381 | bench_fns( |
| 382 | label, |
| 383 | sub_label, |
| 384 | f"marlin ({name_type_string})", |
| 385 | [marlin_create_bench_fn(bt) for bt in benchmark_tensors], |
| 386 | ) |
| 387 | ) |
| 388 | |
| 389 | # machete |
| 390 | timers.append( |
| 391 | bench_fns( |
| 392 | label, |
| 393 | sub_label, |
| 394 | f"machete ({name_type_string})", |
| 395 | [ |
| 396 | machete_create_bench_fn(bt, out_type=types.output_type) |
| 397 | for bt in benchmark_tensors |
| 398 | ], |
| 399 | ) |
| 400 | ) |
| 401 | |
| 402 | # cutlass w4a8 |
| 403 | if types.act_type == torch.float8_e4m3fn and group_size == 128: |
| 404 | timers.append( |
| 405 | bench_fns( |
| 406 | label, |
| 407 | sub_label, |
| 408 | f"cutlass w4a8 ({name_type_string})", |
| 409 | [ |
| 410 | cutlass_w4a8_create_bench_fn(bt, out_type=types.output_type) |
| 411 | for bt in benchmark_tensors |
| 412 | ], |
| 413 | ) |
| 414 | ) |
| 415 | |
| 416 | if sweep_schedules: |
| 417 | global _SWEEP_SCHEDULES_RESULTS |
| 418 | |
| 419 | print("Finding best schedule for machete") |
| 420 | best = None |
| 421 | best_schedule = None |
| 422 | schedules = ops.machete_supported_schedules( |
| 423 | a_type=types.act_type, |
| 424 | b_type=types.weight_type, |
| 425 | group_scales_type=types.group_scale_type, |
| 426 | group_zeros_type=types.group_zero_type, |
| 427 | token_scales_type=types.token_scale_type, |
| 428 | channel_scales_type=types.channel_scale_type, |
| 429 | out_type=types.output_type, |
| 430 | ) |
| 431 | |
| 432 | if schedules is None or len(schedules) == 0: |
| 433 | raise ValueError("No schedules found to sweep") |
| 434 | |
| 435 | for schedule in reversed(schedules): |
| 436 | schedule_M = int(schedule.split("_")[0].split("x")[1]) |
| 437 | |
| 438 | # Prune known bad schedules |
| 439 | if schedule_M >= 2 * max(m, 16) or schedule_M < m // 4: |
| 440 | continue |
| 441 | |
| 442 | res = bench_fns( |
| 443 | label, |
| 444 | sub_label, |
| 445 | "machete_best", |
| 446 | [ |
| 447 | machete_create_bench_fn( |
| 448 | bt, out_type=types.output_type, schedule=schedule |
| 449 | ) |
| 450 | for bt in benchmark_tensors |
| 451 | ], |
| 452 | ) |
| 453 | |
| 454 | results_row = { |
| 455 | "M": m, |
| 456 | "K": k, |
| 457 | "N": n, |
| 458 | "group_size": group_size, |
| 459 | "schedule": schedule, |
| 460 | "median": res.median, |
| 461 | } |
| 462 | if _SWEEP_SCHEDULES_RESULTS is None: |
| 463 | _SWEEP_SCHEDULES_RESULTS = pd.DataFrame(columns=results_row.keys()) |
| 464 | _SWEEP_SCHEDULES_RESULTS.loc[len(_SWEEP_SCHEDULES_RESULTS)] = results_row |
| 465 | |
| 466 | print(f" {res.median:5.5} ", schedule) |
| 467 | if not best or res.median < best.median: |
| 468 | best = res |
| 469 | best_schedule = schedule |
| 470 | print("Best schedule:", best_schedule) |
| 471 | timers.append(best) |
| 472 | |
| 473 | return timers |
| 474 | |
| 475 | |
| 476 | # runner |
| 477 | def print_timers(timers: list[TMeasurement]): |
| 478 | compare = TBenchmark.Compare(timers) |
| 479 | compare.print() |
| 480 | |
| 481 | |
| 482 | def run(args, MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]: |
| 483 | types = TypeConfig( |
| 484 | act_type=args.act_type, |
| 485 | weight_type=scalar_types.uint4b8 |
| 486 | if args.group_zero_type is None |
| 487 | else scalar_types.uint4, |
| 488 | output_type=args.out_type, |
| 489 | group_scale_type=args.group_scale_type, |
| 490 | group_zero_type=args.group_zero_type, |
| 491 | channel_scale_type=args.channel_scale_type, |
| 492 | token_scale_type=args.token_scale_type, |
| 493 | ) |
| 494 | |
| 495 | results: list[TMeasurement] = [] |
| 496 | for m, k, n in MKNs: |
| 497 | timers = bench( |
| 498 | types, |
| 499 | args.group_size, |
| 500 | m, |
| 501 | k, |
| 502 | n, |
| 503 | f"{args.act_type}-gemm", |
| 504 | f"MKN=({m}x{k}x{n})", |
| 505 | sweep_schedules=args.sweep_schedules, |
| 506 | ) |
| 507 | print_timers(timers) |
| 508 | results.extend(timers) |
| 509 | |
| 510 | return results |
| 511 | |
| 512 | |
| 513 | # output makers |
| 514 | def make_output( |
| 515 | data: list[TMeasurement], |
| 516 | MKNs: Iterable[tuple[int, int, int]], |
| 517 | base_description: str, |
| 518 | timestamp=None, |
| 519 | ): |
| 520 | print(f"== All Results {base_description} ====") |
| 521 | print_timers(data) |
| 522 | |
| 523 | # pickle all the results |
| 524 | timestamp = int(time.time()) if timestamp is None else timestamp |
| 525 | with open(f"{base_description}-{timestamp}.pkl", "wb") as f: |
| 526 | pkl.dump(data, f) |
| 527 | |
| 528 | |
| 529 | # argparse runners |
| 530 | |
| 531 | |
| 532 | def run_square_bench(args): |
| 533 | dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment)) |
| 534 | MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes)) |
| 535 | data = run(args.dtype, args.sweep_schedules, MKNs) |
| 536 | |
| 537 | make_output(data, MKNs, f"square_bench-{args.dtype}") |
| 538 | |
| 539 | |
| 540 | def run_range_bench(args): |
| 541 | m_start, k_start, n_start = (int(x) for x in args.dim_start.split(",")) |
| 542 | m_end, k_end, n_end = (int(x) for x in args.dim_end.split(",")) |
| 543 | m_increment, k_increment, n_increment = ( |
| 544 | int(x) for x in args.dim_increment.split(",") |
| 545 | ) |
| 546 | Ms = list(range(m_start, m_end + 1, m_increment)) |
| 547 | Ks = list(range(k_start, k_end + 1, k_increment)) |
| 548 | Ns = list(range(n_start, n_end + 1, n_increment)) |
| 549 | MKNs = list(product(Ms, Ks, Ns)) |
| 550 | |
| 551 | data = run(args.dtype, args.sweep_schedules, MKNs) |
| 552 | |
| 553 | make_output(data, MKNs, f"range_bench-{args.dtype}") |
| 554 | |
| 555 | |
| 556 | def run_model_bench(args): |
| 557 | print("Benchmarking models:") |
| 558 | for i, model in enumerate(args.models): |
| 559 | print(f"[{i}] {model}") |
| 560 | |
| 561 | def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]: |
| 562 | KNs = [] |
| 563 | for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]): |
| 564 | KN[tp_split_dim] = KN[tp_split_dim] // tp_size |
| 565 | KNs.append(KN) |
| 566 | return KNs |
| 567 | |
| 568 | model_bench_data = [] |
| 569 | models_tps = list(itertools.product(args.models, args.tp_sizes)) |
| 570 | for model, tp_size in models_tps: |
| 571 | Ms = args.batch_sizes |
| 572 | KNs = model_shapes(model, tp_size) |
| 573 | MKNs = [] |
| 574 | for m in Ms: |
| 575 | for k, n in KNs: |
| 576 | MKNs.append((m, k, n)) |
| 577 | |
| 578 | data = run(args, MKNs) |
| 579 | model_bench_data.append(data) |
| 580 | |
| 581 | type_string = f"{args.act_type}" |
| 582 | |
| 583 | # Print all results |
| 584 | for data, model_tp in zip(model_bench_data, models_tps): |
| 585 | model, tp_size = model_tp |
| 586 | print(f"== Results {type_string} {model}-TP{tp_size} ====") |
| 587 | print_timers(data) |
| 588 | |
| 589 | timestr = time.strftime("%Y%m%d-%H%M%S") |
| 590 | |
| 591 | all_results = [] |
| 592 | for d in model_bench_data: |
| 593 | all_results.extend(d) |
| 594 | |
| 595 | # pickle all data |
| 596 | with open(f"model_bench-{type_string}-{timestr}.pkl", "wb") as f: |
| 597 | args_dict = vars(args) |
| 598 | args_dict.pop("func") |
| 599 | pkl.dump( |
| 600 | { |
| 601 | "args": args_dict, |
| 602 | "results": all_results, |
| 603 | }, |
| 604 | f, |
| 605 | ) |
| 606 | |
| 607 | |
| 608 | if __name__ == "__main__": |
| 609 | |
| 610 | def to_torch_dtype(dt): |
| 611 | return { |
| 612 | "bfloat16": torch.bfloat16, |
| 613 | "float16": torch.float16, |
| 614 | "int8": torch.int8, |
| 615 | "float8_e4m3fn": torch.float8_e4m3fn, |
| 616 | "int": torch.int, |
| 617 | "float": torch.float, |
| 618 | }[dt] |
| 619 | |
| 620 | class ToTorchDtype(argparse.Action): |
| 621 | def __call__(self, parser, namespace, values, option_string=None): |
| 622 | setattr(namespace, self.dest, to_torch_dtype(values)) |
| 623 | |
| 624 | parser = FlexibleArgumentParser( |
| 625 | description=""" |
| 626 | Benchmark Machete GEMM. |
| 627 | |
| 628 | To run square GEMMs: |
| 629 | python3 ./benchmarks/kernels/benchmark_machete.py --dtype float16 square_bench --dim-start 128 --dim-end 512 --dim-increment 64 |
| 630 | |
| 631 | To run constant N and K and sweep M: |
| 632 | python3 ./benchmarks/kernels/benchmark_machete.py --dtype float16 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384 |
| 633 | |
| 634 | To run dimensions from a model: |
| 635 | python3 ./benchmarks/kernels/benchmark_machete.py --dtype float16 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1 |
| 636 | |
| 637 | Output: |
| 638 | - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs. |
| 639 | """, # noqa: E501 |
| 640 | formatter_class=argparse.RawTextHelpFormatter, |
| 641 | ) |
| 642 | parser.add_argument( |
| 643 | "--act-type", |
| 644 | action=ToTorchDtype, |
| 645 | required=True, |
| 646 | choices=["bfloat16", "float16", "int8", "float8_e4m3fn"], |
| 647 | ) |
| 648 | parser.add_argument( |
| 649 | "--group-scale-type", |
| 650 | action=ToTorchDtype, |
| 651 | choices=["bfloat16", "float16"], |
| 652 | ) |
| 653 | parser.add_argument( |
| 654 | "--group-zero-type", |
| 655 | type=to_torch_dtype, |
| 656 | choices=["bfloat16", "float16"], |
| 657 | ) |
| 658 | parser.add_argument( |
| 659 | "--channel-scale-type", |
| 660 | action=ToTorchDtype, |
| 661 | choices=["float"], |
| 662 | ) |
| 663 | parser.add_argument( |
| 664 | "--token-scale-type", |
| 665 | action=ToTorchDtype, |
| 666 | choices=["float"], |
| 667 | ) |
| 668 | parser.add_argument( |
| 669 | "--out-type", |
| 670 | action=ToTorchDtype, |
| 671 | choices=["bfloat16", "float16"], |
| 672 | ) |
| 673 | parser.add_argument( |
| 674 | "--group-size", |
| 675 | type=int, |
| 676 | help="Available options are ['None', '-1', '128'], default=128", |
| 677 | default=128, |
| 678 | ) |
| 679 | parser.add_argument( |
| 680 | "--sweep-schedules", |
| 681 | action="store_true", |
| 682 | help="Run a sweep over all supported schedules", |
| 683 | ) |
| 684 | parser.add_argument( |
| 685 | "--sweep-csv-out", |
| 686 | help="CSV to store sweep results", |
| 687 | default="sch_sweep_results.csv", |
| 688 | ) |
| 689 | subparsers = parser.add_subparsers(dest="cmd", required=True) |
| 690 | |
| 691 | square_parser = subparsers.add_parser("square_bench") |
| 692 | square_parser.add_argument("--dim-start", type=int, required=True) |
| 693 | square_parser.add_argument("--dim-end", type=int, required=True) |
| 694 | square_parser.add_argument("--dim-increment", type=int, required=True) |
| 695 | square_parser.set_defaults(func=run_square_bench) |
| 696 | |
| 697 | range_parser = subparsers.add_parser("range_bench") |
| 698 | range_parser.add_argument( |
| 699 | "--dim-start", |
| 700 | type=str, |
| 701 | required=True, |
| 702 | help="Start value for M,K,N as common separated list", |
| 703 | ) |
| 704 | range_parser.add_argument( |
| 705 | "--dim-end", |
| 706 | type=str, |
| 707 | required=True, |
| 708 | help="End value (inclusive) for M,K,N as common separated list", |
| 709 | ) |
| 710 | range_parser.add_argument( |
| 711 | "--dim-increment", |
| 712 | type=str, |
| 713 | required=True, |
| 714 | help="Increment value for M,K,N as common separated list", |
| 715 | ) |
| 716 | range_parser.set_defaults(func=run_range_bench) |
| 717 | |
| 718 | model_parser = subparsers.add_parser("model_bench") |
| 719 | model_parser.add_argument( |
| 720 | "--models", |
| 721 | nargs="+", |
| 722 | type=str, |
| 723 | default=DEFAULT_MODELS, |
| 724 | choices=WEIGHT_SHAPES.keys(), |
| 725 | ) |
| 726 | model_parser.add_argument( |
| 727 | "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES |
| 728 | ) |
| 729 | model_parser.add_argument( |
| 730 | "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES |
| 731 | ) |
| 732 | model_parser.set_defaults(func=run_model_bench) |
| 733 | |
| 734 | args = parser.parse_args() |
| 735 | |
| 736 | _SWEEP_SCHEDULES_RESULTS_CSV = args.sweep_csv_out |
| 737 | args.func(args) |
| 738 | |
| 739 | if _SWEEP_SCHEDULES_RESULTS is not None: |
| 740 | _SWEEP_SCHEDULES_RESULTS.to_csv(_SWEEP_SCHEDULES_RESULTS_CSV) |