benchmark_marlin.py 11.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from typing import List
4
5
6
7
8
9

import torch
import torch.utils.benchmark as benchmark
from benchmark_shapes import WEIGHT_SHAPES

from vllm import _custom_ops as ops
10
11
from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
    GPTQ_MARLIN_24_MAX_PARALLEL, GPTQ_MARLIN_24_MIN_THREAD_N,
12
    GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES, GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)
13
14
from vllm.model_executor.layers.quantization.utils.allspark_utils import (
    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, ALLSPARK_SUPPORTED_QUANT_TYPES)
15
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
16
    GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
17
    MARLIN_SUPPORTED_GROUP_SIZES, query_marlin_supported_quant_types)
18
19
20
21
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
    MarlinWorkspace, marlin_quantize)
from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
    marlin_24_quantize)
22
from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
    gptq_pack, gptq_quantize_weights, quantize_weights, sort_weights)
24
from vllm.scalar_type import ScalarType
25
from vllm.utils import FlexibleArgumentParser
26
27

DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"]
28
DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
29
30
31
32
33

ACT_ORDER_OPTS = [False, True]
K_FULL_OPTS = [False, True]


34
def bench_run(results: List[benchmark.Measurement], model: str,
35
36
              act_order: bool, is_k_full: bool, quant_type: ScalarType,
              group_size: int, size_m: int, size_k: int, size_n: int):
37
38
    label = "Quant Matmul"

39
40
41
42
    sub_label = ("{}, act={} k_full={}, q={}, g={}, "
                 "MKN=({}x{}x{})".format(model, act_order, is_k_full,
                                         str(quant_type), group_size, size_m,
                                         size_k, size_n))
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

    print(f"Testing: {sub_label}")

    a = torch.randn(size_m, size_k).to(torch.half).cuda()
    b = torch.rand(size_k, size_n).to(torch.half).cuda()

    a_tmp = (torch.zeros(size_m, size_k).to(torch.half).cuda())

    # Marlin quant
    (
        marlin_w_ref,
        marlin_q_w,
        marlin_s,
        marlin_g_idx,
        marlin_sort_indices,
        marlin_rand_perm,
59
    ) = marlin_quantize(b, quant_type, group_size, act_order)
60

61
62
    # Marlin_24 quant
    (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta,
63
     marlin_24_s) = marlin_24_quantize(b, quant_type, group_size)
64

65
66
    marlin_zp = torch.empty(0, dtype=torch.int, device=b.device)

67
68
    # GPTQ quant
    (w_ref, q_w, s, g_idx,
69
70
     rand_perm) = gptq_quantize_weights(b, quant_type, group_size, act_order)
    q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
71
72
73
74
75
76
77
78

    # For act_order, sort the "weights" and "g_idx"
    # so that group ids are increasing
    repack_sort_indices = torch.empty(0, dtype=torch.int, device=b.device)
    if act_order:
        (q_w, g_idx, repack_sort_indices) = sort_weights(q_w, g_idx)

    # Prepare
79
80
81
82
83
    marlin_workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
                                       GPTQ_MARLIN_MAX_PARALLEL)

    marlin_24_workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_24_MIN_THREAD_N,
                                          GPTQ_MARLIN_24_MAX_PARALLEL)
84
    marlin_zp = torch.zeros_like(marlin_s, dtype=torch.int)
85

86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
    # AllSpark W8A16 quant
    as_supported_case = (quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
                         and group_size == -1 and not act_order and is_k_full)
    if as_supported_case:
        properties = torch.cuda.get_device_properties(b.device.index)
        sm_count = properties.multi_processor_count
        sm_version = properties.major * 10 + properties.minor

        supported_arch = (sm_version >= 80 and sm_version < 90)
        as_supported_case = as_supported_case and supported_arch
        if supported_arch:
            has_zp = False
            w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size,
                                                has_zp)
            qw = qw.to(torch.uint8)

            qw_reorder, s_reorder, zp_reorder = \
                ops.allspark_repack_weight(
                qw, s, zp, has_zp)
            CUBLAS_M_THRESHOLD = ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD

107
    globals = {
108
        # Gen params
109
        "quant_type": quant_type,
110
111
112
113
114
115
116
        "group_size": group_size,
        "size_m": size_m,
        "size_n": size_n,
        "size_k": size_k,
        "a": a,
        "a_tmp": a_tmp,
        # Marlin params
117
118
119
        "marlin_w_ref": marlin_w_ref,
        "marlin_q_w": marlin_q_w,
        "marlin_s": marlin_s,
120
        "marlin_zp": marlin_zp,
121
122
123
        "marlin_g_idx": marlin_g_idx,
        "marlin_sort_indices": marlin_sort_indices,
        "marlin_rand_perm": marlin_rand_perm,
124
125
126
127
128
129
130
131
132
        "marlin_workspace": marlin_workspace,
        "is_k_full": is_k_full,
        # Marlin_24 params
        "marlin_24_w_ref": marlin_24_w_ref,
        "marlin_24_q_w_comp": marlin_24_q_w_comp,
        "marlin_24_meta": marlin_24_meta,
        "marlin_24_s": marlin_24_s,
        "marlin_24_workspace": marlin_24_workspace,
        # GPTQ params
133
134
        "q_w_gptq": q_w_gptq,
        "repack_sort_indices": repack_sort_indices,
135
136
137
138
139
140
141
142
        # AllSpark W8A16 params
        "qw_reorder": qw_reorder if as_supported_case else None,
        "s_reorder": s_reorder if as_supported_case else None,
        "zp_reorder": zp_reorder if as_supported_case else None,
        "sm_count": sm_count if as_supported_case else None,
        "sm_version": sm_version if as_supported_case else None,
        "CUBLAS_M_THRESHOLD":
        CUBLAS_M_THRESHOLD if as_supported_case else None,
143
        # Kernels
144
        "gptq_marlin_gemm": ops.gptq_marlin_gemm,
145
        "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
146
        "gptq_marlin_repack": ops.gptq_marlin_repack,
147
        "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
    }

    min_run_time = 1

    # Warmup pytorch
    for i in range(5):
        torch.matmul(a, marlin_w_ref)

    results.append(
        benchmark.Timer(
            stmt="torch.matmul(a, marlin_w_ref)",
            globals=globals,
            label=label,
            sub_label=sub_label,
            description="pytorch_gemm",
        ).blocked_autorange(min_run_time=min_run_time))

    results.append(
        benchmark.Timer(
            stmt=
168
            "output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)",  # noqa: E501
169
170
171
172
173
174
175
176
177
            globals=globals,
            label=label,
            sub_label=sub_label,
            description="gptq_marlin_gemm_fp16",
        ).blocked_autorange(min_run_time=min_run_time))

    results.append(
        benchmark.Timer(
            stmt=
178
            "output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)",  # noqa: E501
179
180
181
            globals=globals,
            label=label,
            sub_label=sub_label,
182
            description="gptq_marlin_gemm_fp32",
183
184
        ).blocked_autorange(min_run_time=min_run_time))

185
    if (quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES
186
187
188
189
            and group_size in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES):
        results.append(
            benchmark.Timer(
                stmt=
190
                "output = gptq_marlin_24_gemm(a, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s, marlin_24_workspace.scratch, quant_type, size_m, size_n, size_k)",  # noqa: E501
191
192
193
194
195
196
                globals=globals,
                label=label,
                sub_label=sub_label,
                description="gptq_marlin_24_gemm",
            ).blocked_autorange(min_run_time=min_run_time))

197
198
199
    results.append(
        benchmark.Timer(
            stmt=
200
            "q_res = gptq_marlin_repack(q_w_gptq, repack_sort_indices, size_k, size_n, quant_type.size_bits)",  # noqa: E501
201
202
203
204
205
206
            globals=globals,
            label=label,
            sub_label=sub_label,
            description="gptq_marlin_repack",
        ).blocked_autorange(min_run_time=min_run_time))

207
208
209
210
211
212
213
214
215
216
217
    if as_supported_case:
        results.append(
            benchmark.Timer(
                stmt=
                "output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)",  # noqa: E501
                globals=globals,
                label=label,
                sub_label=sub_label,
                description="allspark_w8a16_gemm_fp32",
            ).blocked_autorange(min_run_time=min_run_time))

218
219
220
221
222
223

def main(args):
    print("Benchmarking models:")
    for i, model in enumerate(args.models):
        print(f"[{i}]  {model}")

224
    results: List[benchmark.Measurement] = []
225
226
227
228
229
230
231
232
233
234
235
236
237

    for model in args.models:
        for layer in WEIGHT_SHAPES[model]:
            size_k = layer[0]
            size_n = layer[1]

            if len(args.limit_k) > 0 and size_k not in args.limit_k:
                continue

            if len(args.limit_n) > 0 and size_n not in args.limit_n:
                continue

            for act_order in ACT_ORDER_OPTS:
238
239
240
241
                if len(args.limit_act_order
                       ) > 0 and act_order not in args.limit_act_order:
                    continue

242
                for is_k_full in K_FULL_OPTS:
243
244
245
246
                    if len(args.limit_k_full
                           ) > 0 and is_k_full not in args.limit_k_full:
                        continue

247
248
249
250
                    for quant_type in query_marlin_supported_quant_types(
                            False):
                        if len(args.limit_num_bits) > 0 and \
                            quant_type.size_bits not in args.limit_num_bits:
251
252
                            continue

253
                        for group_size in MARLIN_SUPPORTED_GROUP_SIZES:
254
255
256
257
258
259
260
261
262
263
264
265
266
                            if len(
                                    args.limit_group_size
                            ) > 0 and group_size not in args.limit_group_size:
                                continue

                            # For act_order, the group_size must be less than
                            # size_k
                            if act_order and (group_size == size_k
                                              or group_size == -1):
                                continue

                            for size_m in args.batch_sizes:
                                bench_run(results, model, act_order, is_k_full,
267
268
                                          quant_type, group_size, size_m,
                                          size_k, size_n)
269
270
271
272
273
274

    compare = benchmark.Compare(results)
    compare.print()


# For quick benchmarking use:
275
#   python benchmark_marlin.py --batch-sizes 1 16 32 --limit-k 4096 --limit-n 4096 --limit-group-size 128 --limit-num-bits 4 --limit-act-order 0 --limit-k-full 1 # noqa E501
276
277
#
if __name__ == "__main__":
278
    parser = FlexibleArgumentParser(
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
        description="Benchmark Marlin across specified models/shapes/batches")
    parser.add_argument(
        "--models",
        nargs="+",
        type=str,
        default=DEFAULT_MODELS,
        choices=WEIGHT_SHAPES.keys(),
    )
    parser.add_argument("--batch-sizes",
                        nargs="+",
                        type=int,
                        default=DEFAULT_BATCH_SIZES)
    parser.add_argument("--limit-k", nargs="+", type=int, default=[])
    parser.add_argument("--limit-n", nargs="+", type=int, default=[])
    parser.add_argument("--limit-group-size", nargs="+", type=int, default=[])
294
295
296
    parser.add_argument("--limit-num-bits", nargs="+", type=int, default=[])
    parser.add_argument("--limit-act-order", nargs="+", type=int, default=[])
    parser.add_argument("--limit-k-full", nargs="+", type=int, default=[])
297
298
299

    args = parser.parse_args()
    main(args)