A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 06:00 UTC
[Benchmark] Refactor benchmark script for fp8 & int8 (#19627)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 15, 2025 at 03:15 UTC
3d330c4c095b78b3e6226d99f4d4a7a0965f3758
2 files changed
+184
-280
benchmarks/kernels/bench_fp8_gemm.py
+92
-157
@@ -1,5 +1,4 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
2
import argparse
3
import copy
4
import itertools
@@ -11,6 +10,80 @@ from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
10
from vllm._custom_ops import scaled_fp8_quant as vllm_scaled_fp8_quant
11
from vllm.triton_utils import triton
12
13
+PROVIDER_CFGS = {
14
+ "torch-bf16": dict(enabled=True),
15
+ "fp8-tensor-w-token-a": dict(
16
+ w="tensor", a="token", no_a_quant=False, enabled=False
17
+ ),
18
+ "fp8-tensor-w-tensor-a": dict(
19
+ w="tensor", a="tensor", no_a_quant=False, enabled=True
20
+ ),
21
+ "fp8-channel-w-token-a": dict(
22
+ w="channel", a="token", no_a_quant=False, enabled=True
23
+ ),
24
+ "fp8-channel-w-tensor-a": dict(
25
+ w="channel", a="tensor", no_a_quant=False, enabled=False
26
+ ),
27
+ "fp8-tensor-w-token-a-noquant": dict(
28
+ w="tensor", a="token", no_a_quant=True, enabled=False
29
+ ),
30
+ "fp8-tensor-w-tensor-a-noquant": dict(
31
+ w="tensor", a="tensor", no_a_quant=True, enabled=True
32
+ ),
33
+ "fp8-channel-w-token-a-noquant": dict(
34
+ w="channel", a="token", no_a_quant=True, enabled=True
35
+ ),
36
+ "fp8-channel-w-tensor-a-noquant": dict(
37
+ w="channel", a="tensor", no_a_quant=True, enabled=False
38
+ ),
39
+}
40
+
41
+_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
42
+
43
+
44
+def _quant_weight_fp8(b: torch.Tensor, w_type: str, device: str):
45
+ if w_type == "tensor":
46
+ scale_b = torch.ones(1, device=device, dtype=torch.float32)
47
+ b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
48
+ else:
49
+ b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, use_per_token_if_dynamic=True)
50
+ return b_fp8.t(), scale_b_fp8
51
+
52
+
53
+def build_fp8_runner(cfg, a, b, dtype, device):
54
+ b_fp8, scale_b_fp8 = _quant_weight_fp8(b, cfg["w"], device)
55
+
56
+ scale_a_const = (
57
+ torch.ones(1, device=device, dtype=torch.float32)
58
+ if cfg["a"] == "tensor"
59
+ else None
60
+ )
61
+
62
+ if cfg["no_a_quant"]:
63
+ if cfg["a"] == "tensor":
64
+ a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_const)
65
+ else:
66
+ a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, use_per_token_if_dynamic=True)
67
+
68
+ def run():
69
+ return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
70
+
71
+ return run
72
+
73
+ if cfg["a"] == "tensor":
74
+
75
+ def run():
76
+ a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_const)
77
+ return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
78
+
79
+ else:
80
+
81
+ def run():
82
+ a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, use_per_token_if_dynamic=True)
83
+ return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
84
+
85
+ return run
86
+
87
88
@triton.testing.perf_report(
89
triton.testing.Benchmark(
@@ -18,28 +91,8 @@ from vllm.triton_utils import triton
91
x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
92
x_log=False,
93
line_arg="provider",
21
- line_vals=[
22
- "torch-bf16",
23
- # "fp8-tensor-w-token-a",
24
- "fp8-tensor-w-tensor-a",
25
- "fp8-channel-w-token-a",
26
- # "fp8-channel-w-tensor-a",
27
- # "fp8-tensor-w-token-a-noquant",
28
- "fp8-tensor-w-tensor-a-noquant",
29
- "fp8-channel-w-token-a-noquant",
30
- # "fp8-channel-w-tensor-a-noquant",
31
- ],
32
- line_names=[
33
- "torch-bf16",
34
- # "fp8-tensor-w-token-a",
35
- "fp8-tensor-w-tensor-a",
36
- "fp8-channel-w-token-a",
37
- # "fp8-channel-w-tensor-a",
38
- # "fp8-tensor-w-token-a-noquant",
39
- "fp8-tensor-w-tensor-a-noquant",
40
- "fp8-channel-w-token-a-noquant",
41
- # "fp8-channel-w-tensor-a-noquant",
42
- ],
94
+ line_vals=_enabled,
95
+ line_names=_enabled,
96
ylabel="TFLOP/s (larger is better)",
97
plot_name="BF16 vs FP8 GEMMs",
98
args={},
@@ -50,144 +103,34 @@ def benchmark(batch_size, provider, N, K):
103
device = "cuda"
104
dtype = torch.bfloat16
105
53
- # Create input tensors
106
a = torch.randn((M, K), device=device, dtype=dtype)
107
b = torch.randn((N, K), device=device, dtype=dtype)
108
109
quantiles = [0.5, 0.2, 0.8]
110
59
- if "torch-bf16" in provider:
111
+ if provider == "torch-bf16":
112
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
113
lambda: torch.nn.functional.linear(a, b), quantiles=quantiles
114
)
63
-
64
- elif "fp8" in provider:
65
- # Weights are always quantized ahead of time
66
- if "noquant" in provider:
67
- # For no quantization, we just measure the GEMM
68
- if "tensor-w-token-a" in provider:
69
- # Dynamic per-token quant for A, per-tensor quant for B
70
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b)
71
- assert scale_b_fp8.numel() == 1
72
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(
73
- a, use_per_token_if_dynamic=True
74
- )
75
-
76
- def run_quant():
77
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
78
-
79
- elif "tensor-w-tensor-a" in provider:
80
- # Static per-tensor quantization with fixed scales
81
- # for both A and B
82
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
83
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
84
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
85
- assert scale_b_fp8.numel() == 1
86
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a)
87
-
88
- def run_quant():
89
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
90
-
91
- elif "channel-w-token-a" in provider:
92
- # Static per-channel quantization for weights, per-token
93
- # quant for A
94
- scale_b = torch.tensor((N,), device=device, dtype=torch.float32)
95
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
96
- scale_b_fp8 = scale_b_fp8.expand(N).contiguous()
97
- assert scale_b_fp8.numel() == N
98
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(
99
- a, use_per_token_if_dynamic=True
100
- )
101
-
102
- def run_quant():
103
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
104
-
105
- elif "channel-w-tensor-a" in provider:
106
- # Static per-channel quantization for weights, per-tensor
107
- # quant for A
108
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
109
- scale_b = torch.tensor((N,), device=device, dtype=torch.float32)
110
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
111
- scale_b_fp8 = scale_b_fp8.expand(N).contiguous()
112
- assert scale_b_fp8.numel() == N
113
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a)
114
-
115
- def run_quant():
116
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
117
-
118
- else:
119
- # In these cases, we quantize the activations during the GEMM call
120
- if "tensor-w-token-a" in provider:
121
- # Dynamic per-token quant for A, per-tensor quant for B
122
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b)
123
- assert scale_b_fp8.numel() == 1
124
-
125
- def run_quant():
126
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(
127
- a, use_per_token_if_dynamic=True
128
- )
129
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
130
-
131
- elif "tensor-w-tensor-a" in provider:
132
- # Static per-tensor quantization with fixed scales
133
- # for both A and B
134
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
135
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
136
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
137
- assert scale_b_fp8.numel() == 1
138
-
139
- def run_quant():
140
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a)
141
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
142
-
143
- elif "channel-w-token-a" in provider:
144
- # Static per-channel quantization for weights, per-token
145
- # quant for A
146
- scale_b = torch.tensor((N,), device=device, dtype=torch.float32)
147
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
148
- scale_b_fp8 = scale_b_fp8.expand(N).contiguous()
149
- assert scale_b_fp8.numel() == N
150
-
151
- def run_quant():
152
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(
153
- a, use_per_token_if_dynamic=True
154
- )
155
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
156
-
157
- elif "channel-w-tensor-a" in provider:
158
- # Static per-channel quantization for weights, per-tensor
159
- # quant for A
160
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
161
- scale_b = torch.tensor((N,), device=device, dtype=torch.float32)
162
- b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b)
163
- scale_b_fp8 = scale_b_fp8.expand(N).contiguous()
164
- assert scale_b_fp8.numel() == N
165
-
166
- def run_quant():
167
- a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a)
168
- return vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype)
169
-
170
- b_fp8 = b_fp8.t()
171
-
115
+ else:
116
+ cfg = PROVIDER_CFGS[provider]
117
+ run_quant = build_fp8_runner(cfg, a, b, dtype, device)
118
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
119
lambda: run_quant(), quantiles=quantiles
120
)
121
176
- # Calculate TFLOP/s, two flops per multiply-add
177
- tflops = lambda ms: (2 * M * N * K) * 1e-12 / (ms * 1e-3)
178
- return tflops(ms), tflops(max_ms), tflops(min_ms)
122
+ to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
123
+ return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
124
125
126
def prepare_shapes(args):
182
- KN_model_names = []
183
- models_tps = list(itertools.product(args.models, args.tp_sizes))
184
- for model, tp_size in models_tps:
185
- assert model in WEIGHT_SHAPES
186
- for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
187
- KN[tp_split_dim] = KN[tp_split_dim] // tp_size
127
+ out = []
128
+ for model, tp_size in itertools.product(args.models, args.tp_sizes):
129
+ for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
130
+ KN[tp_dim] //= tp_size
131
KN.append(model)
189
- KN_model_names.append(KN)
190
- return KN_model_names
132
+ out.append(KN)
133
+ return out
134
135
136
if __name__ == "__main__":
@@ -197,21 +140,13 @@ if __name__ == "__main__":
140
nargs="+",
141
type=str,
142
default=["meta-llama/Llama-3.1-8B-Instruct"],
200
- choices=[*WEIGHT_SHAPES.keys()],
201
- help="List of models to benchmark",
202
- )
203
- parser.add_argument(
204
- "--tp-sizes",
205
- nargs="+",
206
- type=int,
207
- default=[1],
208
- help="List of tensor parallel sizes",
143
+ choices=list(WEIGHT_SHAPES.keys()),
144
)
145
+ parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
146
args = parser.parse_args()
147
212
- KN_model_names = prepare_shapes(args)
213
- for K, N, model_name in KN_model_names:
214
- print(f"{model_name}, N={N} K={K}, BF16 vs FP8 GEMMs TFLOP/s:")
148
+ for K, N, model in prepare_shapes(args):
149
+ print(f"{model}, N={N} K={K}, BF16 vs FP8 GEMMs TFLOP/s:")
150
benchmark.run(
151
print_data=True,
152
show_plots=True,
benchmarks/kernels/bench_int8_gemm.py
+92
-123
@@ -11,6 +11,84 @@ 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(
@@ -18,28 +96,8 @@ from vllm.triton_utils import triton
96
x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
97
x_log=False,
98
line_arg="provider",
21
- line_vals=[
22
- "torch-bf16",
23
- # "int8-tensor-w-token-a",
24
- "int8-tensor-w-tensor-a",
25
- "int8-channel-w-token-a",
26
- # "int8-channel-w-tensor-a",
27
- # "int8-tensor-w-token-a-noquant",
28
- "int8-tensor-w-tensor-a-noquant",
29
- "int8-channel-w-token-a-noquant",
30
- # "int8-channel-w-tensor-a-noquant",
31
- ],
32
- line_names=[
33
- "torch-bf16",
34
- # "int8-tensor-w-token-a",
35
- "int8-tensor-w-tensor-a",
36
- "int8-channel-w-token-a",
37
- # "int8-channel-w-tensor-a",
38
- # "int8-tensor-w-token-a-noquant",
39
- "int8-tensor-w-tensor-a-noquant",
40
- "int8-channel-w-token-a-noquant",
41
- # "int8-channel-w-tensor-a-noquant",
42
- ],
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={},
@@ -54,114 +112,26 @@ def benchmark(batch_size, provider, N, K):
112
113
quantiles = [0.5, 0.2, 0.8]
114
57
- if "torch-bf16" in provider:
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
)
61
-
62
- elif "int8" in provider:
63
- # Weights are always quantized ahead of time
64
- if "noquant" in provider:
65
- # For "no quant", we don't measure the time for activations
66
- if "tensor-w-token-a" in provider:
67
- # Dynamic per-token quant for A, static per-tensor quant for B
68
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
69
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
70
- assert scale_b_int8.numel() == 1
71
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
72
-
73
- elif "tensor-w-tensor-a" in provider:
74
- # Static per-tensor quantization with fixed scales for both A and B
75
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
76
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
77
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
78
- assert scale_b_int8.numel() == 1
79
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
80
-
81
- elif "channel-w-token-a" in provider:
82
- # Dynamic per-channel quantization for weights, per-token quant for A
83
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
84
- assert scale_b_int8.numel() == N
85
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
86
-
87
- elif "channel-w-tensor-a" in provider:
88
- # Dynamic per-channel quantization for weights, per-tensor quant for A
89
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
90
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
91
- assert scale_b_int8.numel() == N
92
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
93
-
94
- def run_quant():
95
- return vllm_scaled_mm(a_int8, b_int8, scale_a_int8, scale_b_int8, dtype)
96
-
97
- else:
98
- # Quantize the activations during the GEMM call
99
- if "tensor-w-token-a" in provider:
100
- # Dynamic per-token quant for A, static per-tensor quant for B
101
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
102
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
103
- assert scale_b_int8.numel() == 1
104
-
105
- def run_quant():
106
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
107
- return vllm_scaled_mm(
108
- a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
109
- )
110
-
111
- elif "tensor-w-tensor-a" in provider:
112
- # Static per-tensor quantization with fixed scales for both A and B
113
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
114
- scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
115
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
116
- assert scale_b_int8.numel() == 1
117
-
118
- def run_quant():
119
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
120
- return vllm_scaled_mm(
121
- a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
122
- )
123
-
124
- elif "channel-w-token-a" in provider:
125
- # Dynamic per-channel quant for weights, per-token quant for A
126
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
127
- assert scale_b_int8.numel() == N
128
-
129
- def run_quant():
130
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
131
- return vllm_scaled_mm(
132
- a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
133
- )
134
-
135
- elif "channel-w-tensor-a" in provider:
136
- # Dynamic per-channel quant for weights, static per-tensor quant for A
137
- scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
138
- b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
139
- assert scale_b_int8.numel() == N
140
-
141
- def run_quant():
142
- a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
143
- return vllm_scaled_mm(
144
- a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
145
- )
146
-
147
- b_int8 = b_int8.t()
148
-
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
153
- # Calculate TFLOP/s, two flops per multiply-add
154
- tflops = lambda ms: (2 * M * N * K) * 1e-12 / (ms * 1e-3)
155
- return tflops(ms), tflops(max_ms), tflops(min_ms)
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 = []
160
- models_tps = list(itertools.product(args.models, args.tp_sizes))
161
- for model, tp_size in models_tps:
162
- assert model in WEIGHT_SHAPES
163
- for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
164
- KN[tp_split_dim] = KN[tp_split_dim] // tp_size
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
@@ -174,7 +144,7 @@ if __name__ == "__main__":
144
nargs="+",
145
type=str,
146
default=["meta-llama/Llama-3.1-8B-Instruct"],
177
- choices=[*WEIGHT_SHAPES.keys()],
147
+ choices=list(WEIGHT_SHAPES.keys()),
148
help="List of models to benchmark",
149
)
150
parser.add_argument(
@@ -186,9 +156,8 @@ if __name__ == "__main__":
156
)
157
args = parser.parse_args()
158
189
- KN_model_names = prepare_shapes(args)
190
- for K, N, model_name in KN_model_names:
191
- print(f"{model_name}, N={N} K={K}, BF16 vs INT8 GEMMs TFLOP/s:")
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,