| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import copy |
| 5 | import itertools |
| 6 | |
| 7 | import torch |
| 8 | from weight_shapes import WEIGHT_SHAPES |
| 9 | |
| 10 | from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm |
| 11 | from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant |
| 12 | from vllm.triton_utils import triton |
| 13 | |
| 14 | PROVIDER_CFGS = { |
| 15 | "torch-bf16": dict(enabled=True), |
| 16 | "int8-tensor-w-token-a": dict( |
| 17 | w="tensor", a="token", no_a_quant=False, enabled=False |
| 18 | ), |
| 19 | "int8-tensor-w-tensor-a": dict( |
| 20 | w="tensor", a="tensor", no_a_quant=False, enabled=True |
| 21 | ), |
| 22 | "int8-channel-w-token-a": dict( |
| 23 | w="channel", a="token", no_a_quant=False, enabled=True |
| 24 | ), |
| 25 | "int8-channel-w-tensor-a": dict( |
| 26 | w="channel", a="tensor", no_a_quant=False, enabled=False |
| 27 | ), |
| 28 | "int8-tensor-w-token-a-noquant": dict( |
| 29 | w="tensor", a="token", no_a_quant=True, enabled=False |
| 30 | ), |
| 31 | "int8-tensor-w-tensor-a-noquant": dict( |
| 32 | w="tensor", a="tensor", no_a_quant=True, enabled=True |
| 33 | ), |
| 34 | "int8-channel-w-token-a-noquant": dict( |
| 35 | w="channel", a="token", no_a_quant=True, enabled=True |
| 36 | ), |
| 37 | "int8-channel-w-tensor-a-noquant": dict( |
| 38 | w="channel", a="tensor", no_a_quant=True, enabled=False |
| 39 | ), |
| 40 | } |
| 41 | |
| 42 | |
| 43 | def _quant_weight(b, w_type, device): |
| 44 | if w_type == "tensor": |
| 45 | scale_b = torch.ones(1, device=device, dtype=torch.float32) |
| 46 | b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b) |
| 47 | assert scale_b_int8.numel() == 1 |
| 48 | else: # channel |
| 49 | b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b) |
| 50 | assert scale_b_int8.numel() == b.shape[0] |
| 51 | return b_int8.t(), scale_b_int8 |
| 52 | |
| 53 | |
| 54 | def build_int8_runner(cfg, a, b, dtype, device): |
| 55 | # quant before running the kernel |
| 56 | b_int8, scale_b_int8 = _quant_weight(b, cfg["w"], device) |
| 57 | |
| 58 | scale_a_const = None |
| 59 | if cfg["a"] == "tensor": |
| 60 | scale_a_const = torch.ones(1, device=device, dtype=torch.float32) |
| 61 | |
| 62 | # no quant, create activation ahead |
| 63 | if cfg["no_a_quant"]: |
| 64 | if cfg["a"] == "tensor": |
| 65 | a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a_const) |
| 66 | else: # token |
| 67 | a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a) |
| 68 | |
| 69 | def run_quant(): |
| 70 | return vllm_scaled_mm(a_int8, b_int8, scale_a_int8, scale_b_int8, dtype) |
| 71 | |
| 72 | return run_quant |
| 73 | |
| 74 | # dynamic quant, create activation inside |
| 75 | if cfg["a"] == "tensor": |
| 76 | |
| 77 | def run_quant(): |
| 78 | a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a_const) |
| 79 | return vllm_scaled_mm(a_int8, b_int8, scale_a_int8, scale_b_int8, dtype) |
| 80 | |
| 81 | else: # token |
| 82 | |
| 83 | def run_quant(): |
| 84 | a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a) |
| 85 | return vllm_scaled_mm(a_int8, b_int8, scale_a_int8, scale_b_int8, dtype) |
| 86 | |
| 87 | return run_quant |
| 88 | |
| 89 | |
| 90 | _enabled = [k for k, v in PROVIDER_CFGS.items() if v.get("enabled")] |
| 91 | |
| 92 | |
| 93 | @triton.testing.perf_report( |
| 94 | triton.testing.Benchmark( |
| 95 | x_names=["batch_size"], |
| 96 | x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384], |
| 97 | x_log=False, |
| 98 | line_arg="provider", |
| 99 | line_vals=_enabled, |
| 100 | line_names=[k for k in _enabled], |
| 101 | ylabel="TFLOP/s (larger is better)", |
| 102 | plot_name="BF16 vs INT8 GEMMs", |
| 103 | args={}, |
| 104 | ) |
| 105 | ) |
| 106 | def benchmark(batch_size, provider, N, K): |
| 107 | M = batch_size |
| 108 | device = "cuda" |
| 109 | dtype = torch.bfloat16 |
| 110 | a = torch.randn((M, K), device=device, dtype=dtype) |
| 111 | b = torch.randn((N, K), device=device, dtype=dtype) |
| 112 | |
| 113 | quantiles = [0.5, 0.2, 0.8] |
| 114 | |
| 115 | if provider == "torch-bf16": |
| 116 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 117 | lambda: torch.nn.functional.linear(a, b), quantiles=quantiles |
| 118 | ) |
| 119 | else: |
| 120 | cfg = PROVIDER_CFGS[provider] |
| 121 | run_quant = build_int8_runner(cfg, a, b, dtype, device) |
| 122 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 123 | lambda: run_quant(), quantiles=quantiles |
| 124 | ) |
| 125 | |
| 126 | to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3) |
| 127 | return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms) |
| 128 | |
| 129 | |
| 130 | def prepare_shapes(args): |
| 131 | KN_model_names = [] |
| 132 | for model, tp_size in itertools.product(args.models, args.tp_sizes): |
| 133 | for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]): |
| 134 | KN[tp_dim] //= tp_size |
| 135 | KN.append(model) |
| 136 | KN_model_names.append(KN) |
| 137 | return KN_model_names |
| 138 | |
| 139 | |
| 140 | if __name__ == "__main__": |
| 141 | parser = argparse.ArgumentParser() |
| 142 | parser.add_argument( |
| 143 | "--models", |
| 144 | nargs="+", |
| 145 | type=str, |
| 146 | default=["meta-llama/Llama-3.1-8B-Instruct"], |
| 147 | choices=list(WEIGHT_SHAPES.keys()), |
| 148 | help="List of models to benchmark", |
| 149 | ) |
| 150 | parser.add_argument( |
| 151 | "--tp-sizes", |
| 152 | nargs="+", |
| 153 | type=int, |
| 154 | default=[1], |
| 155 | help="List of tensor parallel sizes", |
| 156 | ) |
| 157 | args = parser.parse_args() |
| 158 | |
| 159 | for K, N, model in prepare_shapes(args): |
| 160 | print(f"{model}, N={N} K={K}, BF16 vs INT8 GEMMs TFLOP/s:") |
| 161 | benchmark.run( |
| 162 | print_data=True, |
| 163 | show_plots=True, |
| 164 | save_path=f"bench_int8_res_n{N}_k{K}", |
| 165 | N=N, |
| 166 | K=K, |
| 167 | ) |
| 168 | |
| 169 | print("Benchmark finished!") |