@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 263 lines 8.67 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark: Fused FP8 output quantization in merge_attn_states.
4
5 Compares fused vs unfused approaches for producing FP8-quantized merged
6 attention output:
7 1. Fused CUDA -- single CUDA kernel (merge + FP8 quant)
8 2. Fused Triton -- single Triton kernel (merge + FP8 quant)
9 3. Unfused CUDA -- CUDA merge + torch.compiled FP8 quant
10 4. Unfused Triton -- Triton merge + torch.compiled FP8 quant
11
12 Usage:
13 python benchmarks/fused_kernels/merge_attn_states_benchmarks.py
14 python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --tp 1 4 8
15 python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --dtype bfloat16
16 """
17
18 import argparse
19 import itertools
20
21 import torch
22
23 from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
24 from vllm.benchmarks.lib.utils import default_vllm_config
25 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
26 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
27 from vllm.platforms import current_platform
28 from vllm.triton_utils import triton
29 from vllm.v1.attention.ops.triton_merge_attn_states import (
30 merge_attn_states as merge_attn_states_triton,
31 )
32
33 # ---------------------------------------------------------------------------
34 # Configuration defaults
35 # ---------------------------------------------------------------------------
36
37 NUM_TOKENS_LIST = [1, 16, 64, 256, 1024, 4096]
38
39 # (label, num_heads, head_size) — num_heads is for TP=1
40 HEAD_CONFIGS = [
41 ("DeepSeek-V3 MLA", 128, 128),
42 ("Llama-70B", 64, 128),
43 ("Llama-8B", 32, 128),
44 ]
45
46 TP_SIZES = [1, 2, 4, 8]
47
48 INPUT_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
49
50 QUANTILES = [0.5, 0.2, 0.8]
51
52
53 # ---------------------------------------------------------------------------
54 # Helpers
55 # ---------------------------------------------------------------------------
56
57
58 def short_dtype(dtype: torch.dtype) -> str:
59 return str(dtype).removeprefix("torch.")
60
61
62 def make_inputs(
63 num_tokens: int,
64 num_heads: int,
65 head_size: int,
66 dtype: torch.dtype,
67 ):
68 """Create random prefix/suffix outputs and LSEs."""
69 prefix_output = torch.randn(
70 (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
71 )
72 suffix_output = torch.randn(
73 (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
74 )
75 prefix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
76 suffix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
77 # Sprinkle some inf values to exercise edge-case paths
78 mask = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
79 prefix_lse[mask] = float("inf")
80 mask2 = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
81 suffix_lse[mask2] = float("inf")
82 return prefix_output, suffix_output, prefix_lse, suffix_lse
83
84
85 def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes):
86 """Build (num_tokens, num_heads, head_size, dtype_str) config tuples,
87 applying TP division to num_heads and skipping invalid combos."""
88 configs = []
89 for (_, nh, hs), nt, dtype, tp in itertools.product(
90 head_configs, num_tokens_list, input_dtypes, tp_sizes
91 ):
92 nh_tp = nh // tp
93 if nh_tp >= 1:
94 configs.append((nt, nh_tp, hs, short_dtype(dtype)))
95 return configs
96
97
98 def parse_args():
99 parser = argparse.ArgumentParser(
100 description="Benchmark merge_attn_states fused FP8 quantization"
101 )
102 parser.add_argument(
103 "--num-tokens",
104 type=int,
105 nargs="+",
106 default=None,
107 help=f"Override token counts (default: {NUM_TOKENS_LIST})",
108 )
109 parser.add_argument(
110 "--tp",
111 type=int,
112 nargs="+",
113 default=None,
114 help=f"TP sizes to simulate (divides num_heads) (default: {TP_SIZES})",
115 )
116 parser.add_argument(
117 "--dtype",
118 type=str,
119 nargs="+",
120 default=None,
121 help="Input dtypes (e.g. bfloat16 float16 float32). "
122 f"Default: {[short_dtype(d) for d in INPUT_DTYPES]}",
123 )
124 return parser.parse_args()
125
126
127 # ---------------------------------------------------------------------------
128 # Parse args and build configs before decorators
129 # ---------------------------------------------------------------------------
130
131 args = parse_args()
132
133 num_tokens_list = args.num_tokens if args.num_tokens else NUM_TOKENS_LIST
134 tp_sizes = args.tp if args.tp else TP_SIZES
135
136 if args.dtype:
137 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
138
139 input_dtypes = [STR_DTYPE_TO_TORCH_DTYPE[d] for d in args.dtype]
140 else:
141 input_dtypes = INPUT_DTYPES
142
143 configs = build_configs(HEAD_CONFIGS, num_tokens_list, input_dtypes, tp_sizes)
144
145 torch._dynamo.config.recompile_limit = 8888
146
147
148 # ---------------------------------------------------------------------------
149 # Benchmark function
150 # ---------------------------------------------------------------------------
151
152
153 @triton.testing.perf_report(
154 triton.testing.Benchmark(
155 x_names=["num_tokens", "num_heads", "head_size", "dtype_str"],
156 x_vals=configs,
157 line_arg="provider",
158 line_vals=["fused_cuda", "fused_triton", "unfused_cuda", "unfused_triton"],
159 line_names=["Fused CUDA", "Fused Triton", "Unfused CUDA", "Unfused Triton"],
160 styles=[("blue", "-"), ("green", "-"), ("blue", "--"), ("green", "--")],
161 ylabel="us",
162 plot_name="merge_attn_states FP8 (fused vs unfused)",
163 args={},
164 )
165 )
166 @default_vllm_config()
167 def benchmark(num_tokens, num_heads, head_size, dtype_str, provider):
168 input_dtype = getattr(torch, dtype_str)
169 fp8_dtype = current_platform.fp8_dtype()
170 prefix_out, suffix_out, prefix_lse, suffix_lse = make_inputs(
171 num_tokens, num_heads, head_size, input_dtype
172 )
173 output_scale = torch.tensor([0.1], dtype=torch.float32, device="cuda")
174
175 if provider == "fused_cuda":
176 output = torch.empty(
177 (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
178 )
179 fn = lambda: merge_attn_states_cuda(
180 output,
181 prefix_out,
182 prefix_lse,
183 suffix_out,
184 suffix_lse,
185 output_scale=output_scale,
186 )
187 elif provider == "fused_triton":
188 output = torch.empty(
189 (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
190 )
191 fn = lambda: merge_attn_states_triton(
192 output,
193 prefix_out,
194 prefix_lse,
195 suffix_out,
196 suffix_lse,
197 output_scale=output_scale,
198 )
199 elif provider == "unfused_cuda":
200 merge_buf = torch.empty(
201 (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
202 )
203 quant_fp8 = QuantFP8(
204 static=True,
205 group_shape=GroupShape.PER_TENSOR,
206 column_major_scales=False,
207 )
208 quant_input = merge_buf.view(-1, head_size)
209 compiled_quant = torch.compile(
210 quant_fp8.forward_native, fullgraph=True, dynamic=False
211 )
212
213 def unfused_fn():
214 merge_attn_states_cuda(
215 merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
216 )
217 compiled_quant(quant_input, output_scale)
218
219 fn = unfused_fn
220 else: # unfused_triton
221 merge_buf = torch.empty(
222 (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
223 )
224 quant_fp8 = QuantFP8(
225 static=True,
226 group_shape=GroupShape.PER_TENSOR,
227 column_major_scales=False,
228 )
229 quant_input = merge_buf.view(-1, head_size)
230 compiled_quant = torch.compile(
231 quant_fp8.forward_native, fullgraph=True, dynamic=False
232 )
233
234 def unfused_fn():
235 merge_attn_states_triton(
236 merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
237 )
238 compiled_quant(quant_input, output_scale)
239
240 fn = unfused_fn
241
242 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=QUANTILES)
243 return 1000 * ms, 1000 * max_ms, 1000 * min_ms # us
244
245
246 # ---------------------------------------------------------------------------
247 # Main
248 # ---------------------------------------------------------------------------
249
250
251 def main():
252 device_name = current_platform.get_device_name()
253 print(f"Device: {device_name}")
254 print(f"Token counts: {num_tokens_list}")
255 print(f"TP sizes: {tp_sizes}")
256 print(f"Input dtypes: {[short_dtype(d) for d in input_dtypes]}")
257 print(f"Head configs: {[(c[0], c[1], c[2]) for c in HEAD_CONFIGS]}")
258 benchmark.run(print_data=True)
259
260
261 if __name__ == "__main__":
262 with torch.inference_mode():
263 main()