@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 211 lines 5.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable, Iterable
5 from dataclasses import dataclass
6 from itertools import product
7
8 import torch
9 import torch.nn.functional as F
10 import torch.utils.benchmark as TBenchmark
11 from torch.utils.benchmark import Measurement as TMeasurement
12 from tqdm import tqdm
13
14 import vllm._custom_ops as ops
15 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
16 per_token_group_quant_fp8,
17 )
18
19
20 @dataclass
21 class bench_params_t:
22 num_tokens: int
23 hidden_size: int
24 dtype: torch.dtype
25 group_size: int # Changed from list[int] to int
26
27 def description(self):
28 return (
29 f"N {self.num_tokens} "
30 f"x D {self.hidden_size} "
31 f"x DT {self.dtype} "
32 f"x GS {self.group_size}"
33 )
34
35
36 def get_bench_params() -> list[bench_params_t]:
37 """Test configurations covering common model sizes."""
38 NUM_TOKENS = [16, 128, 512, 2048]
39 HIDDEN_SIZES = [1024, 2048, 4096, 5120, 14336] # Common FFN sizes
40 DTYPES = [torch.float16, torch.bfloat16]
41 GROUP_SIZES = [64, 128] # Changed from [[1, 64], [1, 128]]
42
43 combinations = product(NUM_TOKENS, HIDDEN_SIZES, DTYPES, GROUP_SIZES)
44 bench_params = list(
45 map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations)
46 )
47 return bench_params
48
49
50 # Reference implementations
51 def unfused_fp8_impl(
52 x: torch.Tensor,
53 quant_dtype: torch.dtype,
54 group_size: int, # Changed from list[int]
55 ):
56 """Unfused: SiLU+Mul then per-tensor quantize."""
57 hidden = x.shape[-1] // 2
58 gate, up = x.split(hidden, dim=-1)
59
60 # SiLU(gate) * up
61 silu_out = F.silu(gate) * up
62
63 # Per-tensor quantize (no group_size used here)
64 silu_out, _ = ops.scaled_fp8_quant(silu_out)
65
66
67 def unfused_groupwise_fp8_impl(
68 x: torch.Tensor,
69 quant_dtype: torch.dtype,
70 group_size: int, # Changed from list[int]
71 ):
72 """Unfused: SiLU+Mul then group-wise quantize."""
73 hidden = x.shape[-1] // 2
74 gate, up = x.split(hidden, dim=-1)
75
76 # SiLU(gate) * up
77 silu_out = F.silu(gate) * up
78
79 # Group quantize - use group_size directly
80 silu_out, _ = per_token_group_quant_fp8(
81 silu_out, group_size=group_size, use_ue8m0=False
82 )
83
84
85 def fused_impl(
86 x: torch.Tensor,
87 quant_dtype: torch.dtype,
88 group_size: int,
89 ):
90 """Fused: SiLU+Mul+Block Quantization in single kernel."""
91 out, _ = ops.silu_and_mul_per_block_quant(
92 x,
93 group_size=group_size,
94 quant_dtype=quant_dtype,
95 is_scale_transposed=False,
96 )
97
98
99 # Bench functions
100 def bench_fn(
101 x: torch.Tensor,
102 quant_dtype: torch.dtype,
103 group_size: int,
104 label: str,
105 sub_label: str,
106 fn: Callable,
107 description: str,
108 ) -> TMeasurement:
109 min_run_time = 1
110
111 globals = {
112 "x": x,
113 "quant_dtype": quant_dtype,
114 "group_size": group_size,
115 "fn": fn,
116 }
117 return TBenchmark.Timer(
118 stmt="fn(x, quant_dtype, group_size)",
119 globals=globals,
120 label=label,
121 sub_label=sub_label,
122 description=description,
123 ).blocked_autorange(min_run_time=min_run_time)
124
125
126 def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasurement]:
127 """Run benchmarks for all implementations."""
128 # Make inputs: [num_tokens, hidden_size * 2] for [gate || up]
129 scale = 1 / params.hidden_size
130 x = (
131 torch.randn(
132 params.num_tokens,
133 params.hidden_size * 2,
134 dtype=params.dtype,
135 device="cuda",
136 )
137 * scale
138 )
139
140 timers = []
141
142 # Unfused per-tensor FP8
143 timers.append(
144 bench_fn(
145 x,
146 torch.float8_e4m3fn,
147 params.group_size,
148 label,
149 sub_label,
150 unfused_fp8_impl,
151 "unfused_fp8_impl",
152 )
153 )
154
155 # Unfused group-wise FP8
156 timers.append(
157 bench_fn(
158 x,
159 torch.float8_e4m3fn,
160 params.group_size,
161 label,
162 sub_label,
163 unfused_groupwise_fp8_impl,
164 "unfused_groupwise_fp8_impl",
165 )
166 )
167
168 # Fused group-wise FP8
169 timers.append(
170 bench_fn(
171 x,
172 torch.float8_e4m3fn,
173 params.group_size,
174 label,
175 sub_label,
176 fused_impl,
177 "fused_groupwise_fp8_impl",
178 )
179 )
180
181 return timers
182
183
184 def print_timers(timers: Iterable[TMeasurement]):
185 compare = TBenchmark.Compare(timers)
186 compare.print()
187
188
189 def main():
190 torch.set_default_device("cuda")
191 bench_params = get_bench_params()
192
193 print(f"Running {len(bench_params)} benchmark configurations...")
194 print(
195 f"This will take approximately {len(bench_params) * 3} seconds (1s per variant)"
196 )
197 print()
198
199 timers = []
200 for bp in tqdm(bench_params):
201 result_timers = bench(bp, "silu-mul-block-quant", bp.description())
202 timers.extend(result_timers)
203
204 print("\n" + "=" * 80)
205 print("FINAL COMPARISON - ALL RESULTS")
206 print("=" * 80)
207 print_timers(timers)
208
209
210 if __name__ == "__main__":
211 main()