| 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_fp8_quant as vllm_scaled_fp8_quant |
| 12 | from vllm.triton_utils import triton |
| 13 | |
| 14 | PROVIDER_CFGS = { |
| 15 | "torch-bf16": dict(enabled=True), |
| 16 | "fp8-tensor-w-token-a": dict( |
| 17 | w="tensor", a="token", no_a_quant=False, enabled=False |
| 18 | ), |
| 19 | "fp8-tensor-w-tensor-a": dict( |
| 20 | w="tensor", a="tensor", no_a_quant=False, enabled=True |
| 21 | ), |
| 22 | "fp8-channel-w-token-a": dict( |
| 23 | w="channel", a="token", no_a_quant=False, enabled=True |
| 24 | ), |
| 25 | "fp8-channel-w-tensor-a": dict( |
| 26 | w="channel", a="tensor", no_a_quant=False, enabled=False |
| 27 | ), |
| 28 | "fp8-tensor-w-token-a-noquant": dict( |
| 29 | w="tensor", a="token", no_a_quant=True, enabled=False |
| 30 | ), |
| 31 | "fp8-tensor-w-tensor-a-noquant": dict( |
| 32 | w="tensor", a="tensor", no_a_quant=True, enabled=True |
| 33 | ), |
| 34 | "fp8-channel-w-token-a-noquant": dict( |
| 35 | w="channel", a="token", no_a_quant=True, enabled=True |
| 36 | ), |
| 37 | "fp8-channel-w-tensor-a-noquant": dict( |
| 38 | w="channel", a="tensor", no_a_quant=True, enabled=False |
| 39 | ), |
| 40 | } |
| 41 | |
| 42 | _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]] |
| 43 | |
| 44 | |
| 45 | def _quant_weight_fp8(b: torch.Tensor, w_type: str, device: str): |
| 46 | if w_type == "tensor": |
| 47 | scale_b = torch.ones(1, device=device, dtype=torch.float32) |
| 48 | b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b) |
| 49 | else: |
| 50 | b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, use_per_token_if_dynamic=True) |
| 51 | return b_fp8.t(), scale_b_fp8 |
| 52 | |
| 53 | |
| 54 | def build_fp8_runner(cfg, a, b, dtype, device): |
| 55 | b_fp8, scale_b_fp8 = _quant_weight_fp8(b, cfg["w"], device) |
| 56 | |
| 57 | scale_a_const = ( |
| 58 | torch.ones(1, device=device, dtype=torch.float32) |
| 59 | if cfg["a"] == "tensor" |
| 60 | else None |
| 61 | ) |
| 62 | |
| 63 | if cfg["no_a_quant"]: |
| 64 | if cfg["a"] == "tensor": |
| 65 | a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_const) |
| 66 | else: |
| 67 | a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, use_per_token_if_dynamic=True) |
| 68 | |
| 69 | def run(): |
| 70 | return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype) |
| 71 | |
| 72 | return run |
| 73 | |
| 74 | if cfg["a"] == "tensor": |
| 75 | |
| 76 | def run(): |
| 77 | a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_const) |
| 78 | return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype) |
| 79 | |
| 80 | else: |
| 81 | |
| 82 | def run(): |
| 83 | a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, use_per_token_if_dynamic=True) |
| 84 | return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype) |
| 85 | |
| 86 | return run |
| 87 | |
| 88 | |
| 89 | @triton.testing.perf_report( |
| 90 | triton.testing.Benchmark( |
| 91 | x_names=["batch_size"], |
| 92 | x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384], |
| 93 | x_log=False, |
| 94 | line_arg="provider", |
| 95 | line_vals=_enabled, |
| 96 | line_names=_enabled, |
| 97 | ylabel="TFLOP/s (larger is better)", |
| 98 | plot_name="BF16 vs FP8 GEMMs", |
| 99 | args={}, |
| 100 | ) |
| 101 | ) |
| 102 | def benchmark(batch_size, provider, N, K): |
| 103 | M = batch_size |
| 104 | device = "cuda" |
| 105 | dtype = torch.bfloat16 |
| 106 | |
| 107 | a = torch.randn((M, K), device=device, dtype=dtype) |
| 108 | b = torch.randn((N, K), device=device, dtype=dtype) |
| 109 | |
| 110 | quantiles = [0.5, 0.2, 0.8] |
| 111 | |
| 112 | if provider == "torch-bf16": |
| 113 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 114 | lambda: torch.nn.functional.linear(a, b), quantiles=quantiles |
| 115 | ) |
| 116 | else: |
| 117 | cfg = PROVIDER_CFGS[provider] |
| 118 | run_quant = build_fp8_runner(cfg, a, b, dtype, device) |
| 119 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 120 | lambda: run_quant(), quantiles=quantiles |
| 121 | ) |
| 122 | |
| 123 | to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3) |
| 124 | return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms) |
| 125 | |
| 126 | |
| 127 | def prepare_shapes(args): |
| 128 | out = [] |
| 129 | for model, tp_size in itertools.product(args.models, args.tp_sizes): |
| 130 | for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]): |
| 131 | KN[tp_dim] //= tp_size |
| 132 | KN.append(model) |
| 133 | out.append(KN) |
| 134 | return out |
| 135 | |
| 136 | |
| 137 | if __name__ == "__main__": |
| 138 | parser = argparse.ArgumentParser() |
| 139 | parser.add_argument( |
| 140 | "--models", |
| 141 | nargs="+", |
| 142 | type=str, |
| 143 | default=["meta-llama/Llama-3.1-8B-Instruct"], |
| 144 | choices=list(WEIGHT_SHAPES.keys()), |
| 145 | ) |
| 146 | parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1]) |
| 147 | args = parser.parse_args() |
| 148 | |
| 149 | for K, N, model in prepare_shapes(args): |
| 150 | print(f"{model}, N={N} K={K}, BF16 vs FP8 GEMMs TFLOP/s:") |
| 151 | benchmark.run( |
| 152 | print_data=True, |
| 153 | show_plots=True, |
| 154 | save_path=f"bench_fp8_res_n{N}_k{K}", |
| 155 | N=N, |
| 156 | K=K, |
| 157 | ) |
| 158 | |
| 159 | print("Benchmark finished!") |