| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.utils.benchmark as benchmark |
| 6 | from benchmark_shapes import WEIGHT_SHAPES |
| 7 | |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 10 | GPTQ_MARLIN_MAX_PARALLEL, |
| 11 | GPTQ_MARLIN_MIN_THREAD_N, |
| 12 | MARLIN_SUPPORTED_GROUP_SIZES, |
| 13 | query_marlin_supported_quant_types, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 16 | FP4_MARLIN_SUPPORTED_GROUP_SIZES, |
| 17 | rand_marlin_weight_fp4_like, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 20 | marlin_quant_fp8_torch, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( |
| 23 | MarlinWorkspace, |
| 24 | awq_marlin_quantize, |
| 25 | marlin_quantize, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 28 | gptq_pack, |
| 29 | gptq_quantize_weights, |
| 30 | ) |
| 31 | from vllm.scalar_type import ScalarType, scalar_types |
| 32 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 33 | |
| 34 | DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"] |
| 35 | DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 36 | |
| 37 | |
| 38 | def bench_run( |
| 39 | results: list[benchmark.Measurement], |
| 40 | model: str, |
| 41 | quant_type: ScalarType, |
| 42 | group_size: int, |
| 43 | size_m: int, |
| 44 | size_k: int, |
| 45 | size_n: int, |
| 46 | ): |
| 47 | label = "Quant Matmul" |
| 48 | sub_label = "{}, q={}, g={}, MKN=({}x{}x{})".format( |
| 49 | model, str(quant_type), group_size, size_m, size_k, size_n |
| 50 | ) |
| 51 | print(f"Testing: {sub_label}") |
| 52 | |
| 53 | a = torch.randn(size_m, size_k).to(torch.half).cuda() |
| 54 | b = torch.rand(size_k, size_n).to(torch.half).cuda() |
| 55 | has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] |
| 56 | if size_k % group_size != 0: |
| 57 | return |
| 58 | |
| 59 | repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES |
| 60 | |
| 61 | def gen_marlin_params(): |
| 62 | # Marlin quant |
| 63 | marlin_zp = marlin_s2 = None |
| 64 | if quant_type == scalar_types.float4_e2m1f: |
| 65 | if group_size != 16: |
| 66 | return |
| 67 | marlin_w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_fp4_like( |
| 68 | b.T, group_size |
| 69 | ) |
| 70 | elif quant_type == scalar_types.float8_e4m3fn: |
| 71 | if group_size not in [-1, 128]: |
| 72 | return |
| 73 | marlin_w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(b.T, group_size) |
| 74 | elif group_size == 16: |
| 75 | return |
| 76 | elif has_zp: |
| 77 | marlin_w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize( |
| 78 | b, quant_type, group_size |
| 79 | ) |
| 80 | else: |
| 81 | marlin_w_ref, marlin_q_w, marlin_s = marlin_quantize( |
| 82 | b, quant_type, group_size |
| 83 | ) |
| 84 | return ( |
| 85 | marlin_w_ref, |
| 86 | marlin_q_w, |
| 87 | marlin_s, |
| 88 | marlin_s2, |
| 89 | marlin_zp, |
| 90 | ) |
| 91 | |
| 92 | def gen_repack_params(): |
| 93 | q_w_gptq = None |
| 94 | if repack_supported: |
| 95 | _, q_w, _ = gptq_quantize_weights(b, quant_type, group_size) |
| 96 | q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n) |
| 97 | return q_w_gptq |
| 98 | |
| 99 | ( |
| 100 | marlin_w_ref, |
| 101 | marlin_q_w, |
| 102 | marlin_s, |
| 103 | marlin_s2, |
| 104 | marlin_zp, |
| 105 | ) = gen_marlin_params() |
| 106 | q_w_gptq = gen_repack_params() |
| 107 | |
| 108 | # Prepare |
| 109 | marlin_workspace = MarlinWorkspace( |
| 110 | size_n, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL |
| 111 | ) |
| 112 | |
| 113 | globals = { |
| 114 | # Gen params |
| 115 | "quant_type": quant_type, |
| 116 | "group_size": group_size, |
| 117 | "size_m": size_m, |
| 118 | "size_n": size_n, |
| 119 | "size_k": size_k, |
| 120 | "a": a, |
| 121 | # Marlin params |
| 122 | "marlin_w_ref": marlin_w_ref, |
| 123 | "marlin_q_w": marlin_q_w, |
| 124 | "marlin_s": marlin_s, |
| 125 | "marlin_s2": marlin_s2, |
| 126 | "marlin_zp": marlin_zp, |
| 127 | "marlin_workspace": marlin_workspace, |
| 128 | # GPTQ params |
| 129 | "q_w_gptq": q_w_gptq, |
| 130 | # Kernels |
| 131 | "marlin_gemm": ops.marlin_gemm, |
| 132 | "gptq_marlin_repack": ops.gptq_marlin_repack, |
| 133 | } |
| 134 | |
| 135 | min_run_time = 1 |
| 136 | |
| 137 | # Warmup pytorch |
| 138 | for _ in range(5): |
| 139 | torch.matmul(a, marlin_w_ref) |
| 140 | |
| 141 | results.append( |
| 142 | benchmark.Timer( |
| 143 | stmt="torch.matmul(a, marlin_w_ref)", |
| 144 | globals=globals, |
| 145 | label=label, |
| 146 | sub_label=sub_label, |
| 147 | description="pytorch_gemm", |
| 148 | ).blocked_autorange(min_run_time=min_run_time) |
| 149 | ) |
| 150 | |
| 151 | results.append( |
| 152 | benchmark.Timer( |
| 153 | stmt="output = marlin_gemm(a, None, marlin_q_w, None, marlin_s, None, marlin_s2, marlin_zp, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, False, False, False)", # noqa: E501 |
| 154 | globals=globals, |
| 155 | label=label, |
| 156 | sub_label=sub_label, |
| 157 | description="marlin_gemm", |
| 158 | ).blocked_autorange(min_run_time=min_run_time) |
| 159 | ) |
| 160 | |
| 161 | results.append( |
| 162 | benchmark.Timer( |
| 163 | stmt="output = marlin_gemm(a, None, marlin_q_w, None, marlin_s, None, marlin_s2, marlin_zp, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, False, True, False)", # noqa: E501 |
| 164 | globals=globals, |
| 165 | label=label, |
| 166 | sub_label=sub_label, |
| 167 | description="marlin_gemm_fp32", |
| 168 | ).blocked_autorange(min_run_time=min_run_time) |
| 169 | ) |
| 170 | |
| 171 | if repack_supported: |
| 172 | results.append( |
| 173 | benchmark.Timer( |
| 174 | stmt="q_res = gptq_marlin_repack(q_w_gptq, size_k, size_n, quant_type.size_bits)", # noqa: E501 |
| 175 | globals=globals, |
| 176 | label=label, |
| 177 | sub_label=sub_label, |
| 178 | description="gptq_marlin_repack", |
| 179 | ).blocked_autorange(min_run_time=min_run_time) |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | def main(args): |
| 184 | print("Benchmarking models:") |
| 185 | for i, model in enumerate(args.models): |
| 186 | print(f"[{i}] {model}") |
| 187 | results: list[benchmark.Measurement] = [] |
| 188 | |
| 189 | for model in args.models: |
| 190 | for layer in WEIGHT_SHAPES[model]: |
| 191 | size_k = layer[0] |
| 192 | size_n = layer[1] |
| 193 | |
| 194 | if len(args.limit_k) > 0 and size_k not in args.limit_k: |
| 195 | continue |
| 196 | |
| 197 | if len(args.limit_n) > 0 and size_n not in args.limit_n: |
| 198 | continue |
| 199 | |
| 200 | for quant_type in query_marlin_supported_quant_types(): |
| 201 | if ( |
| 202 | len(args.limit_num_bits) > 0 |
| 203 | and quant_type.size_bits not in args.limit_num_bits |
| 204 | ): |
| 205 | continue |
| 206 | |
| 207 | for group_size in ( |
| 208 | MARLIN_SUPPORTED_GROUP_SIZES + FP4_MARLIN_SUPPORTED_GROUP_SIZES |
| 209 | ): |
| 210 | if ( |
| 211 | len(args.limit_group_size) > 0 |
| 212 | and group_size not in args.limit_group_size |
| 213 | ): |
| 214 | continue |
| 215 | |
| 216 | for size_m in args.batch_sizes: |
| 217 | bench_run( |
| 218 | results, |
| 219 | model, |
| 220 | quant_type, |
| 221 | group_size, |
| 222 | size_m, |
| 223 | size_k, |
| 224 | size_n, |
| 225 | ) |
| 226 | |
| 227 | compare = benchmark.Compare(results) |
| 228 | compare.print() |
| 229 | |
| 230 | |
| 231 | # For quick benchmarking use: |
| 232 | # python benchmark_marlin.py --batch-sizes 1 16 32 --limit-k 4096 --limit-n 4096 --limit-group-size 128 --limit-num-bits 4 # noqa E501 |
| 233 | # |
| 234 | if __name__ == "__main__": |
| 235 | parser = FlexibleArgumentParser( |
| 236 | description="Benchmark Marlin across specified models/shapes/batches" |
| 237 | ) |
| 238 | parser.add_argument( |
| 239 | "--models", |
| 240 | nargs="+", |
| 241 | type=str, |
| 242 | default=DEFAULT_MODELS, |
| 243 | choices=WEIGHT_SHAPES.keys(), |
| 244 | ) |
| 245 | parser.add_argument( |
| 246 | "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES |
| 247 | ) |
| 248 | parser.add_argument("--limit-k", nargs="+", type=int, default=[]) |
| 249 | parser.add_argument("--limit-n", nargs="+", type=int, default=[]) |
| 250 | parser.add_argument("--limit-group-size", nargs="+", type=int, default=[]) |
| 251 | parser.add_argument("--limit-num-bits", nargs="+", type=int, default=[]) |
| 252 | |
| 253 | args = parser.parse_args() |
| 254 | main(args) |