@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 451 lines 13.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # Adapted from sglang quantization/tuning_block_wise_kernel.py
4
5 import argparse
6 import json
7 import multiprocessing as mp
8 import os
9 import time
10 from datetime import datetime
11 from typing import Any
12
13 import torch
14 from tqdm import tqdm
15
16 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
17 _w8a8_triton_block_scaled_mm,
18 )
19 from vllm.platforms import current_platform
20 from vllm.triton_utils import triton
21 from vllm.utils.argparse_utils import FlexibleArgumentParser
22 from vllm.utils.platform_utils import get_device_name_as_file_name
23
24 mp.set_start_method("spawn", force=True)
25
26 assert (
27 current_platform.is_cuda()
28 or current_platform.is_rocm()
29 or current_platform.is_xpu()
30 ), "Only support tune w8a8 block fp8 kernel on CUDA/ROCm/XPU device."
31
32 DTYPE_MAP = {
33 "float32": torch.float32,
34 "float16": torch.float16,
35 "half": torch.half,
36 "bfloat16": torch.bfloat16,
37 }
38
39
40 def w8a8_block_matmul(
41 A: torch.Tensor,
42 B: torch.Tensor,
43 As: torch.Tensor,
44 Bs: torch.Tensor,
45 block_size: list[int],
46 config: dict[str, Any],
47 output_dtype: torch.dtype = torch.float16,
48 ) -> torch.Tensor:
49 """This function performs matrix multiplication with
50 block-wise quantization.
51
52 It takes two input tensors `A` and `B` with scales `As` and `Bs`.
53 The output is returned in the specified `output_dtype`.
54
55 Args:
56 A: The input tensor, e.g., activation.
57 B: The input tensor, e.g., weight.
58 As: The per-token-group quantization scale for `A`.
59 Bs: The per-block quantization scale for `B`.
60 block_size: The block size for per-block quantization.
61 It should be 2-dim, e.g., [128, 128].
62 config: Triton kernel launch config, e.g. tile sizes and num_warps.
63 output_dtype: The dtype of the returned tensor.
64
65 Returns:
66 torch.Tensor: The result of matmul.
67
68 """
69 assert len(block_size) == 2
70 block_n, block_k = block_size[0], block_size[1]
71
72 assert A.shape[-1] == B.shape[-1]
73 assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
74 assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
75 M = A.numel() // A.shape[-1]
76
77 assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
78 N, K = B.shape
79 assert triton.cdiv(N, block_n) == Bs.shape[0]
80 assert triton.cdiv(K, block_k) == Bs.shape[1]
81
82 C_shape = A.shape[:-1] + (N,)
83 C = A.new_empty(C_shape, dtype=output_dtype)
84
85 def grid(META):
86 return (
87 triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
88 )
89
90 if A.dtype == torch.float8_e4m3fn:
91 kernel = _w8a8_triton_block_scaled_mm
92 else:
93 raise RuntimeError("Currently, only support tune w8a8 block fp8 kernel.")
94
95 kernel[grid](
96 A,
97 B,
98 C,
99 As,
100 Bs,
101 M,
102 N,
103 K,
104 block_n,
105 block_k,
106 A.stride(-2),
107 A.stride(-1),
108 B.stride(1),
109 B.stride(0),
110 C.stride(-2),
111 C.stride(-1),
112 As.stride(-2),
113 As.stride(-1),
114 Bs.stride(1),
115 Bs.stride(0),
116 **config,
117 )
118
119 return C
120
121
122 def get_configs_compute_bound():
123 configs = []
124 for num_stages in [2, 3, 4, 5]:
125 for block_m in [16, 32, 64, 128, 256]:
126 for block_k in [64, 128]:
127 for block_n in [32, 64, 128, 256]:
128 for num_warps in [4, 8]:
129 for group_size in [1, 16, 32, 64]:
130 configs.append(
131 {
132 "BLOCK_SIZE_M": block_m,
133 "BLOCK_SIZE_N": block_n,
134 "BLOCK_SIZE_K": block_k,
135 "GROUP_SIZE_M": group_size,
136 "num_warps": num_warps,
137 "num_stages": num_stages,
138 }
139 )
140 return configs
141
142
143 def get_xpu_tuning_space():
144 # NOTE: curated subset of get_configs_compute_bound(); the full grid is
145 # too slow to tune on XPU due to per-config Triton JIT compile time.
146 configs = []
147 for num_stages in [2, 3]:
148 for block_m in [32, 64, 128]:
149 for block_k in [64, 128]:
150 for block_n in [128, 256]:
151 for num_warps in [4, 8, 16, 32]:
152 for group_size in [1, 32]:
153 configs.append(
154 {
155 "BLOCK_SIZE_M": block_m,
156 "BLOCK_SIZE_N": block_n,
157 "BLOCK_SIZE_K": block_k,
158 "GROUP_SIZE_M": group_size,
159 "num_warps": num_warps,
160 "num_stages": num_stages,
161 }
162 )
163 return configs
164
165
166 def get_weight_shapes(tp_size):
167 # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3.
168 # Modify them, if you tune for another different model.
169 # cannot TP
170 total = [
171 (512 + 64, 7168),
172 (2112, 7168),
173 ((128 + 64) * 128, 7168),
174 (128 * (128 + 128), 512),
175 (7168, 16384),
176 (7168, 18432),
177 ]
178 # N can TP
179 n_tp = [
180 (18432 * 2, 7168),
181 ((128 + 64) * 128, 7168),
182 (128 * (128 + 128), 512),
183 (24576, 1536),
184 (12288, 7168),
185 (4096, 7168),
186 ]
187 # K can TP
188 k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)]
189
190 weight_shapes = []
191 for t in total:
192 weight_shapes.append(t)
193 for n_t in n_tp:
194 new_t = (n_t[0] // tp_size, n_t[1])
195 weight_shapes.append(new_t)
196 for k_t in k_tp:
197 new_t = (k_t[0], k_t[1] // tp_size)
198 weight_shapes.append(new_t)
199 return weight_shapes
200
201
202 def benchmark_config(
203 A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10
204 ):
205 def run():
206 w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)
207
208 torch.accelerator.synchronize()
209 # JIT complication & warmup
210 for _ in range(5):
211 run()
212 torch.accelerator.synchronize()
213
214 start_event = torch.Event(enable_timing=True)
215 end_event = torch.Event(enable_timing=True)
216
217 latencies: list[float] = []
218 for i in range(num_iters):
219 torch.accelerator.synchronize()
220 start_event.record()
221 run()
222 end_event.record()
223 end_event.synchronize()
224 latencies.append(start_event.elapsed_time(end_event))
225 avg = sum(latencies) / (num_iters * 10) * 1000 # us
226 return avg
227
228
229 def tune(M, N, K, block_size, out_dtype, search_space, input_type):
230 factor_for_scale = 1e-2
231
232 if input_type == "fp8":
233 fp8_info = torch.finfo(torch.float8_e4m3fn)
234 fp8_max, fp8_min = fp8_info.max, fp8_info.min
235
236 device = current_platform.device_type
237 A_fp32 = (
238 (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
239 )
240 A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
241
242 B_fp32 = (
243 (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
244 )
245 B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
246 else:
247 raise RuntimeError("Currently, only support tune w8a8 block fp8 kernel.")
248
249 block_n, block_k = block_size[0], block_size[1]
250 n_tiles = (N + block_n - 1) // block_n
251 k_tiles = (K + block_k - 1) // block_k
252
253 As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
254 Bs = (
255 torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
256 * factor_for_scale
257 )
258
259 best_config = None
260 best_time = float("inf")
261 for config in tqdm(search_space):
262 try:
263 kernel_time = benchmark_config(
264 A,
265 B,
266 As,
267 Bs,
268 block_size,
269 config,
270 out_dtype,
271 num_iters=10,
272 )
273 except triton.runtime.autotuner.OutOfResources:
274 # Some configurations may be invalid and fail to compile.
275 continue
276
277 if kernel_time < best_time:
278 best_time = kernel_time
279 best_config = config
280 now = datetime.now()
281 print(f"{now.ctime()}] Completed tuning for batch_size={M}")
282 assert best_config is not None
283 return best_config
284
285
286 def save_configs(
287 N,
288 K,
289 block_n,
290 block_k,
291 configs,
292 save_path,
293 input_type="fp8",
294 ) -> None:
295 os.makedirs(save_path, exist_ok=True)
296 device_name = get_device_name_as_file_name()
297 json_file_name = (
298 f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8,"
299 f"block_shape=[{block_n},{block_k}].json"
300 )
301
302 config_file_path = os.path.join(save_path, json_file_name)
303 print(f"Writing best config to {config_file_path}...")
304
305 with open(config_file_path, "w") as f:
306 json.dump(configs, f, indent=4)
307 f.write("\n")
308
309
310 def tune_on_gpu(args_dict):
311 """Run tuning on a specific GPU."""
312 gpu_id = args_dict["gpu_id"]
313 batch_sizes = args_dict["batch_sizes"]
314 weight_shapes = args_dict["weight_shapes"]
315 args = args_dict["args"]
316
317 torch.accelerator.set_device_index(gpu_id)
318 print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}")
319
320 block_n = args.block_n
321 block_k = args.block_k
322 out_dtype = DTYPE_MAP[args.out_dtype]
323 save_path = args.save_path
324 input_type = args.input_type
325
326 search_space = (
327 get_xpu_tuning_space()
328 if current_platform.is_xpu()
329 else get_configs_compute_bound()
330 )
331 search_space = [
332 config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0
333 ]
334
335 start = time.time()
336 for shape in tqdm(weight_shapes, desc=f"GPU {gpu_id} - Shapes"):
337 N, K = shape[0], shape[1]
338 print(f"[GPU {gpu_id}] Tune for weight shape of `N: {N}, K: {K}`")
339 benchmark_results = [
340 tune(
341 batch_size,
342 N,
343 K,
344 [block_n, block_k],
345 out_dtype,
346 search_space,
347 input_type,
348 )
349 for batch_size in tqdm(batch_sizes, desc=f"GPU {gpu_id} - Batch sizes")
350 ]
351 best_configs = {M: config for M, config in zip(batch_sizes, benchmark_results)}
352 save_configs(N, K, block_n, block_k, best_configs, save_path, input_type)
353
354 end = time.time()
355 print(f"Tuning on GPU {gpu_id} took {end - start:.2f} seconds")
356
357
358 def distribute_batch_sizes(batch_sizes, num_gpus):
359 """Distribute batch sizes across available GPUs."""
360 batches_per_gpu = []
361 for i in range(num_gpus):
362 start_idx = i * len(batch_sizes) // num_gpus
363 end_idx = (i + 1) * len(batch_sizes) // num_gpus
364 batches_per_gpu.append(batch_sizes[start_idx:end_idx])
365 return batches_per_gpu
366
367
368 def main(args):
369 print(args)
370 num_gpus = torch.accelerator.device_count()
371 if num_gpus == 0:
372 raise RuntimeError("No GPU available for tuning")
373 print(f"Found {num_gpus} GPUs for parallel tuning")
374
375 if current_platform.is_xpu():
376 torch.xpu.init()
377 else:
378 torch.cuda.init()
379
380 if args.batch_size is None:
381 batch_sizes = [
382 1,
383 2,
384 4,
385 8,
386 16,
387 24,
388 32,
389 48,
390 64,
391 96,
392 128,
393 256,
394 512,
395 1024,
396 1536,
397 2048,
398 3072,
399 4096,
400 ]
401 else:
402 batch_sizes = [args.batch_size]
403 num_gpus = 1 # If only one batch size, use only one GPU
404
405 weight_shapes = get_weight_shapes(args.tp_size)
406
407 batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus)
408
409 process_args = []
410 for gpu_id in range(num_gpus):
411 process_args.append(
412 {
413 "gpu_id": gpu_id,
414 "batch_sizes": batches_per_gpu[gpu_id],
415 "weight_shapes": weight_shapes, # Each GPU processes all weight shapes
416 "args": args,
417 }
418 )
419
420 ctx = mp.get_context("spawn")
421 with ctx.Pool(num_gpus) as pool:
422 pool.map(tune_on_gpu, process_args)
423
424 print("Multi-GPU tuning completed")
425
426
427 if __name__ == "__main__":
428 parser = FlexibleArgumentParser(
429 description="""
430 Tune triton w8a8 block fp8 for DeepSeek-V3/DeepSeek-R1:
431 python3 benchmark_w8a8_block_fp8.py --tp-size 8 --input-type fp8
432 Then copy to model_executor/layers/quantization/utils/configs
433 """,
434 formatter_class=argparse.RawTextHelpFormatter,
435 )
436
437 parser.add_argument("--tp-size", "-tp", type=int, default=8)
438 parser.add_argument("--input-type", type=str, choices=["fp8"], default="fp8")
439 parser.add_argument(
440 "--out-dtype",
441 type=str,
442 choices=["float32", "float16", "bfloat16", "half"],
443 default="float16",
444 )
445 parser.add_argument("--block-n", type=int, default=128)
446 parser.add_argument("--block-k", type=int, default=128)
447 parser.add_argument("--batch-size", type=int, required=False)
448 parser.add_argument("--save-path", type=str, default="./")
449 args = parser.parse_args()
450
451 main(args)