@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 312 lines 7.26 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pickle as pkl
5 import time
6 from collections.abc import Callable, Iterable
7 from dataclasses import dataclass
8 from itertools import product
9
10 import torch
11 import torch.utils.benchmark as TBenchmark
12 from torch.utils.benchmark import Measurement as TMeasurement
13 from tqdm import tqdm
14
15 import vllm._custom_ops as ops
16 from vllm.benchmarks.lib.utils import default_vllm_config
17 from vllm.model_executor.layers.layernorm import RMSNorm
18 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19 per_token_group_quant_fp8,
20 )
21
22
23 @dataclass
24 class bench_params_t:
25 num_tokens: int
26 hidden_size: int
27 add_residual: bool
28 dtype: torch.dtype
29 group_size: list[int]
30
31 def description(self):
32 return (
33 f"N {self.num_tokens} "
34 f"x D {self.hidden_size} "
35 f"x R {self.add_residual} "
36 f"x DT {self.dtype}"
37 f"x GS {self.group_size}"
38 )
39
40
41 def get_bench_params() -> list[bench_params_t]:
42 ## Test Fixtures
43 NUM_TOKENS = [2**x for x in range(11)]
44 HIDDEN_SIZES = list(range(1024, 8129, 1024))
45 ADD_RESIDUAL = [True, False]
46 DTYPES = [torch.bfloat16, torch.float]
47 GROUP_SIZES = [[1, 64], [1, 128]]
48
49 combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES, GROUP_SIZES)
50 bench_params = list(
51 map(lambda x: bench_params_t(x[0], x[1], x[2], x[3], x[4]), combinations)
52 )
53 return bench_params
54
55
56 # Reference impls
57 def unfused_int8_impl(
58 rms_norm_layer: RMSNorm,
59 x: torch.Tensor,
60 residual: torch.Tensor | None,
61 quant_dtype: torch.dtype,
62 group_size: list[int],
63 ):
64 # Norm
65 torch_out = None
66 if residual is None:
67 torch_out = rms_norm_layer.forward_cuda(x, residual)
68 else:
69 torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
70
71 # Quant
72 torch_out, _, _ = ops.scaled_int8_quant(torch_out)
73
74
75 def unfused_fp8_impl(
76 rms_norm_layer: RMSNorm,
77 x: torch.Tensor,
78 residual: torch.Tensor | None,
79 quant_dtype: torch.dtype,
80 group_size: list[int],
81 ):
82 # Norm
83 torch_out = None
84 if residual is None:
85 torch_out = rms_norm_layer.forward_cuda(x, residual)
86 else:
87 torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
88
89 # Quant
90 torch_out, _ = ops.scaled_fp8_quant(torch_out)
91
92
93 def unfused_groupwise_fp8_impl(
94 rms_norm_layer: RMSNorm,
95 x: torch.Tensor,
96 residual: torch.Tensor | None,
97 quant_dtype: torch.dtype,
98 group_size: list[int],
99 ):
100 # Norm
101 torch_out = None
102 if residual is None:
103 torch_out = rms_norm_layer.forward_cuda(x, residual)
104 else:
105 torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
106
107 # Quant
108 torch_out, _ = per_token_group_quant_fp8(
109 torch_out, group_size=group_size[1], use_ue8m0=False
110 )
111
112
113 def fused_impl(
114 rms_norm_layer: RMSNorm, # this stores the weights
115 x: torch.Tensor,
116 residual: torch.Tensor | None,
117 quant_dtype: torch.dtype,
118 group_size: list[int],
119 ):
120 out, _ = ops.rms_norm_dynamic_per_token_quant(
121 x, rms_norm_layer.weight, 1e-6, quant_dtype, residual=residual
122 )
123
124
125 def fused_groupwise_impl(
126 rms_norm_layer: RMSNorm, # this stores the weights
127 x: torch.Tensor,
128 residual: torch.Tensor | None,
129 quant_dtype: torch.dtype,
130 group_size: list[int],
131 ):
132 out, _ = ops.rms_norm_per_block_quant(
133 x,
134 rms_norm_layer.weight,
135 1e-6,
136 quant_dtype,
137 group_size,
138 residual=residual,
139 is_scale_transposed=True,
140 )
141
142
143 # Bench functions
144 def bench_fn(
145 rms_norm_layer: RMSNorm,
146 x: torch.Tensor,
147 residual: torch.Tensor,
148 quant_dtype: torch.dtype,
149 group_size: list[int],
150 label: str,
151 sub_label: str,
152 fn: Callable,
153 description: str,
154 ) -> TMeasurement:
155 min_run_time = 1
156
157 globals = {
158 "rms_norm_layer": rms_norm_layer,
159 "x": x,
160 "residual": residual,
161 "quant_dtype": quant_dtype,
162 "group_size": group_size,
163 "fn": fn,
164 }
165 return TBenchmark.Timer(
166 stmt="fn(rms_norm_layer, x, residual, quant_dtype, group_size)",
167 globals=globals,
168 label=label,
169 sub_label=sub_label,
170 description=description,
171 ).blocked_autorange(min_run_time=min_run_time)
172
173
174 def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasurement]:
175 # Make inputs
176 layer = RMSNorm(params.hidden_size, 1e-6).to(dtype=params.dtype)
177 # Make weights
178 layer.weight.data.normal_(mean=1.0, std=0.1)
179 # Make inputs
180 scale = 1 / params.hidden_size
181 x = (
182 torch.randn(
183 params.num_tokens, params.hidden_size, dtype=params.dtype, device="cuda"
184 )
185 * scale
186 )
187 residual = (
188 (torch.randn_like(x) * scale).to(device="cuda") if params.add_residual else None
189 )
190
191 timers = []
192
193 # unfused int8 impl.
194 timers.append(
195 bench_fn(
196 layer,
197 x,
198 residual,
199 torch.int8,
200 params.group_size,
201 label,
202 sub_label,
203 unfused_int8_impl,
204 "unfused_int8_impl",
205 )
206 )
207
208 # unfused fp8 impl.
209 timers.append(
210 bench_fn(
211 layer,
212 x,
213 residual,
214 torch.float8_e4m3fn,
215 params.group_size,
216 label,
217 sub_label,
218 unfused_fp8_impl,
219 "unfused_fp8_impl",
220 )
221 )
222
223 # fused int8 impl.
224 timers.append(
225 bench_fn(
226 layer,
227 x,
228 residual,
229 torch.int8,
230 params.group_size,
231 label,
232 sub_label,
233 fused_impl,
234 "fused_int8_impl",
235 )
236 )
237
238 # fused fp8 impl.
239 timers.append(
240 bench_fn(
241 layer,
242 x,
243 residual,
244 torch.float8_e4m3fn,
245 params.group_size,
246 label,
247 sub_label,
248 fused_impl,
249 "fused_fp8_impl",
250 )
251 )
252
253 # unfused groupwise fp8 impl.
254 timers.append(
255 bench_fn(
256 layer,
257 x,
258 residual,
259 torch.float8_e4m3fn,
260 params.group_size,
261 label,
262 sub_label,
263 unfused_groupwise_fp8_impl,
264 "unfused_groupwise_fp8_impl",
265 )
266 )
267
268 # fused groupwise fp8 impl.
269 timers.append(
270 bench_fn(
271 layer,
272 x,
273 residual,
274 torch.float8_e4m3fn,
275 params.group_size,
276 label,
277 sub_label,
278 fused_groupwise_impl,
279 "fused_groupwise_fp8_impl",
280 )
281 )
282
283 print_timers(timers)
284
285 return timers
286
287
288 # launch bench
289 # runner
290 def print_timers(timers: Iterable[TMeasurement]):
291 compare = TBenchmark.Compare(timers)
292 compare.print()
293
294
295 @default_vllm_config()
296 def main():
297 torch.set_default_device("cuda")
298 bench_params = get_bench_params()
299
300 timers = []
301 for bp in tqdm(bench_params):
302 timers.extend(bench(bp, "rms-norm-dynamic-per-token-quant", bp.description()))
303 print_timers(timers)
304
305 # pickle all the results
306 timestamp = int(time.time())
307 with open(f"rms_norm_dpt_quant-{timestamp}.pkl", "wb") as f:
308 pkl.dump(timers, f)
309
310
311 if __name__ == "__main__":
312 main()