@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 243 lines 6.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 from enum import Enum
6 from itertools import product
7 from typing import Any
8
9 import torch
10 import torch.utils.benchmark as TBenchmark
11 from torch.utils.benchmark import Measurement as TMeasurement
12
13 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14 _per_token_group_quant_fp8_colmajor,
15 silu_mul_per_token_group_quant_fp8_colmajor,
16 )
17 from vllm.triton_utils import triton
18 from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
19
20 from .utils import ArgPool, Bench, CudaGraphBenchParams
21
22 GROUP_SIZE = 128
23 FLOAT8_T = torch.float8_e4m3fn
24
25
26 def print_timers(timers: list[TMeasurement], cuda_graph_nops: int):
27 print(
28 f"Note : The timings reported above is for {cuda_graph_nops} "
29 "consecutive invocations of the benchmarking functions. "
30 f"Please divide by {cuda_graph_nops} for single invocation "
31 "timings."
32 )
33 compare = TBenchmark.Compare(timers)
34 compare.print()
35
36
37 class ImplType(Enum):
38 SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR = 1
39 REFERENCE = 2
40
41 def get_impl(self):
42 if self == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
43 return silu_mul_per_token_group_quant_fp8_colmajor
44 elif self == ImplType.REFERENCE:
45 return reference
46 raise ValueError(f"Unrecognized ImplType {self}")
47
48
49 @dataclass
50 class BenchmarkTensors:
51 input: torch.Tensor
52 output: torch.Tensor
53
54 # Reference act output tensor
55 ref_act_out: torch.Tensor
56 ref_quant_out: torch.Tensor
57
58 @staticmethod
59 def make(T: int, N: int) -> "BenchmarkTensors":
60 assert T % GROUP_SIZE == 0
61 assert N % (GROUP_SIZE * 2) == 0
62
63 input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda")
64
65 # silu_mul_per_token_group_quant_fp8_colmajor output.
66 output = torch.rand((T, N // 2), dtype=torch.bfloat16, device="cuda").to(
67 FLOAT8_T
68 )
69
70 # reference output.
71 ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda")
72 ref_quant_out = torch.empty(
73 (T, N // 2), dtype=torch.bfloat16, device="cuda"
74 ).to(FLOAT8_T)
75
76 return BenchmarkTensors(
77 input=input,
78 output=output,
79 ref_act_out=ref_act_out,
80 ref_quant_out=ref_quant_out,
81 )
82
83 @property
84 def T(self):
85 return self.input.size(0)
86
87 @property
88 def N(self):
89 return self.input.size(1)
90
91 def make_impl_kwargs(self, impl_type: ImplType) -> dict[str, Any]:
92 if impl_type == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
93 return {
94 "input": self.input,
95 "output": self.output,
96 "use_ue8m0": is_deep_gemm_e8m0_used(),
97 }
98 elif impl_type == ImplType.REFERENCE:
99 return {
100 "input": self.input,
101 "act_out": self.ref_act_out,
102 "quant_out": self.ref_quant_out,
103 "use_ue8m0": is_deep_gemm_e8m0_used(),
104 }
105 raise ValueError(f"Unrecognized impl_type {impl_type}")
106
107
108 def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool):
109 """Reference triton quant kernel from,
110 vllm.model_executor.layers.quantization.utils.fp8_utils
111 """
112 assert quant_out.size() == x.size()
113 # Allocate the scale tensor column-major format.
114 shape = (x.shape[-1] // GROUP_SIZE,) + x.shape[:-1]
115 x_q = quant_out
116 x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)
117
118 M = x.numel() // GROUP_SIZE
119 N = GROUP_SIZE
120 BLOCK = triton.next_power_of_2(N)
121 # heuristics for number of warps
122 num_warps = min(max(BLOCK // 256, 1), 8)
123 num_stages = 1
124
125 finfo = torch.finfo(FLOAT8_T)
126 fp8_min = finfo.min
127 fp8_max = finfo.max
128
129 _per_token_group_quant_fp8_colmajor[(M,)](
130 x,
131 x_q,
132 x_s,
133 GROUP_SIZE,
134 x.shape[1],
135 x.stride(0),
136 x_s.stride(1),
137 eps=1e-10,
138 fp8_min=fp8_min,
139 fp8_max=fp8_max,
140 use_ue8m0=use_ue8m0,
141 BLOCK=BLOCK,
142 num_warps=num_warps,
143 num_stages=num_stages,
144 )
145 return x_q, x_s
146
147
148 def reference(
149 input: torch.Tensor,
150 act_out: torch.Tensor,
151 quant_out: torch.Tensor,
152 use_ue8m0: bool,
153 ) -> tuple[torch.Tensor, torch.Tensor]:
154 torch.ops._C.silu_and_mul(act_out, input)
155 return reference_quant(act_out, quant_out, use_ue8m0)
156
157
158 def bench_impl(
159 bench_tensors: list[BenchmarkTensors], impl_type: ImplType
160 ) -> TMeasurement:
161 T = bench_tensors[0].T
162 N = bench_tensors[0].N
163
164 arg_pool_size = len(bench_tensors)
165 kwargs_list = [bt.make_impl_kwargs(impl_type) for bt in bench_tensors]
166
167 # warmup
168 for kwargs in kwargs_list:
169 impl_type.get_impl()(**kwargs)
170 torch.accelerator.synchronize()
171
172 # Merge into a single kwargs and qualify arguments as ArgPool
173 kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
174 for _kwargs in kwargs_list:
175 for k, v in _kwargs.items():
176 kwargs[k].values.append(v)
177
178 cuda_graph_params = None
179 cuda_graph_params = CudaGraphBenchParams(arg_pool_size)
180 timer = None
181 with Bench(
182 cuda_graph_params,
183 "silu-mul-quant",
184 f"num_tokens={T}, N={N}",
185 impl_type.name,
186 impl_type.get_impl(),
187 **kwargs,
188 ) as bench:
189 timer = bench.run()
190 return timer
191
192
193 def test_correctness(T: int, N: int):
194 print(f"Testing num_tokens={T}, N={N} ...")
195
196 bench_tensor = BenchmarkTensors.make(T, N)
197
198 def output_from_impl(impl: ImplType) -> tuple[torch.Tensor, torch.Tensor]:
199 return impl.get_impl()(**bench_tensor.make_impl_kwargs(impl))
200
201 # reference output
202 ref_out_q, ref_out_s = output_from_impl(ImplType.REFERENCE)
203
204 # test output
205 out_q, out_s = output_from_impl(
206 ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
207 )
208
209 torch.testing.assert_close(ref_out_q.to(torch.float32), out_q.to(torch.float32))
210 torch.testing.assert_close(ref_out_s, out_s)
211
212
213 def run(Ts: list[int], Ns: list[int], arg_pool_size: int) -> list[TMeasurement]:
214 timers = []
215 for N, T in product(Ns, Ts):
216 test_correctness(T, N)
217
218 bench_tensors: list[BenchmarkTensors] = [
219 BenchmarkTensors.make(T, N) for _ in range(arg_pool_size)
220 ]
221
222 silu_mul_quant_timer = bench_impl(
223 bench_tensors, ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
224 )
225 timers.append(silu_mul_quant_timer)
226 reference_timer = bench_impl(bench_tensors, ImplType.REFERENCE)
227 timers.append(reference_timer)
228
229 print_timers(
230 [silu_mul_quant_timer, reference_timer], cuda_graph_nops=arg_pool_size
231 )
232
233 print_timers(timers, cuda_graph_nops=arg_pool_size)
234
235 return timers
236
237
238 if __name__ == "__main__":
239 T = [128 * i for i in range(1, 16)] + [2048 * i for i in range(1, 65)]
240 N = [2048, 4096, 8192]
241
242 print(f"T = {T}, N = {N}")
243 run(T, N, arg_pool_size=8)