@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 167 lines 5.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5
6 # Disable DeepGEMM for this benchmark to use CUTLASS
7 os.environ["VLLM_USE_DEEP_GEMM"] = "0"
8
9 import torch
10
11 from vllm.benchmarks.lib.utils import default_vllm_config
12 from vllm.model_executor.kernels.linear import (
13 init_fp8_linear_kernel,
14 )
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 GroupShape,
17 create_fp8_quant_key,
18 )
19 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20 CUTLASS_BLOCK_FP8_SUPPORTED,
21 )
22 from vllm.platforms import current_platform
23 from vllm.triton_utils import triton as vllm_triton
24
25 assert current_platform.is_cuda(), (
26 "Only support benchmarking w8a8 block fp8 kernel on CUDA device."
27 )
28
29 # DeepSeek-V3 weight shapes
30 DEEPSEEK_V3_SHAPES = [
31 (512 + 64, 7168),
32 (2112, 7168),
33 ((128 + 64) * 128, 7168),
34 (128 * (128 + 128), 512),
35 (7168, 16384),
36 (7168, 18432),
37 (18432 * 2, 7168),
38 (24576, 1536),
39 (12288, 7168),
40 (4096, 7168),
41 (7168, 2048),
42 ]
43
44
45 @default_vllm_config()
46 def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
47 """Build runner function for w8a8 block fp8 matmul."""
48 factor_for_scale = 1e-2
49
50 fp8_info = torch.finfo(torch.float8_e4m3fn)
51 fp8_max, fp8_min = fp8_info.max, fp8_info.min
52
53 # Create random input tensor (bfloat16, will be quantized by W8A8BlockFp8LinearOp)
54 A_ref = (torch.rand(M, K, dtype=torch.bfloat16, device=device) - 0.5) * 2 * fp8_max
55
56 # Create quantized weight tensor
57 B_ref = (torch.rand(N, K, dtype=torch.bfloat16, device=device) - 0.5) * 2 * fp8_max
58 B = B_ref.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
59
60 # Create weight scales
61 block_n, block_k = block_size[0], block_size[1]
62 n_tiles = (N + block_n - 1) // block_n
63 k_tiles = (K + block_k - 1) // block_k
64
65 Bs = (
66 torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
67 * factor_for_scale
68 )
69
70 # Create W8A8BlockFp8LinearOp instance
71 weight_group_shape = GroupShape(block_n, block_k)
72 act_quant_group_shape = GroupShape(1, block_k) # Per-token, per-group quantization
73
74 linear_op = init_fp8_linear_kernel(
75 weight_quant_key=create_fp8_quant_key(
76 static=True, group_shape=weight_group_shape
77 ),
78 activation_quant_key=create_fp8_quant_key(
79 static=False, group_shape=act_quant_group_shape
80 ),
81 out_dtype=torch.get_default_dtype(),
82 module_name="build_w8a8_block_fp8_runner",
83 )
84
85 def run():
86 return linear_op.apply(
87 input=A_ref,
88 weight=B,
89 weight_scale=Bs,
90 input_scale=None,
91 bias=None,
92 )
93
94 return run
95
96
97 # Determine available providers
98 available_providers = ["torch-bf16", "w8a8-block-fp8-triton"]
99 plot_title = "BF16 vs W8A8 Block FP8 GEMMs"
100
101 if CUTLASS_BLOCK_FP8_SUPPORTED:
102 available_providers.append("w8a8-block-fp8-cutlass")
103
104
105 @vllm_triton.testing.perf_report(
106 vllm_triton.testing.Benchmark(
107 x_names=["batch_size"],
108 x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
109 x_log=False,
110 line_arg="provider",
111 line_vals=available_providers,
112 line_names=available_providers,
113 ylabel="TFLOP/s (larger is better)",
114 plot_name="BF16 vs W8A8 Block FP8 GEMMs",
115 args={},
116 )
117 )
118 def benchmark_tflops(batch_size, provider, N, K, block_size=(128, 128)):
119 M = batch_size
120 device = "cuda"
121
122 quantiles = [0.5, 0.2, 0.8]
123
124 if provider == "torch-bf16":
125 a = torch.randn((M, K), device=device, dtype=torch.bfloat16)
126 b = torch.randn((N, K), device=device, dtype=torch.bfloat16)
127 ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
128 lambda: torch.nn.functional.linear(a, b), quantiles=quantiles
129 )
130 elif provider == "w8a8-block-fp8-triton":
131 run_w8a8_triton = build_w8a8_block_fp8_runner(
132 M, N, K, block_size, device, use_cutlass=False
133 )
134 ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
135 lambda: run_w8a8_triton(), quantiles=quantiles
136 )
137 elif provider == "w8a8-block-fp8-cutlass":
138 run_w8a8_cutlass = build_w8a8_block_fp8_runner(
139 M, N, K, block_size, device, use_cutlass=True
140 )
141 ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
142 lambda: run_w8a8_cutlass(), quantiles=quantiles
143 )
144 else:
145 raise ValueError(f"Unknown provider: {provider}")
146
147 to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
148 return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
149
150
151 if __name__ == "__main__":
152 block_size = (128, 128)
153
154 for N, K in DEEPSEEK_V3_SHAPES:
155 print(f"\nBenchmarking DeepSeek-V3, N={N} K={K}")
156
157 print(f"TFLOP/s comparison (block_size={block_size}):")
158 benchmark_tflops.run(
159 print_data=True,
160 # show_plots=False,
161 # save_path=f"bench_w8a8_block_fp8_tflops_n{N}_k{K}",
162 N=N,
163 K=K,
164 block_size=block_size,
165 )
166
167 print("\nBenchmark finished!")