@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 200 lines 6.19 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
4
5 Usage:
6 python benchmark_int4_gemm.py
7 python benchmark_int4_gemm.py --models Qwen/Qwen3-4B
8 python benchmark_int4_gemm.py --group-size 128
9 """
10
11 import argparse
12 import copy
13 import itertools
14 import os
15
16 import torch
17
18 from vllm.triton_utils import triton
19
20 # ---------------------------------------------------------------------------
21 # Weight shapes: [K, N], TP_SPLIT_DIM
22 # ---------------------------------------------------------------------------
23 WEIGHT_SHAPES = {
24 "Qwen/Qwen3-4B": [
25 ([2560, 3840], 1), # qkv_proj
26 ([2560, 2560], 0), # o_proj
27 ([2560, 19456], 1), # gate_up_proj
28 ([9728, 2560], 0), # down_proj
29 ],
30 "Qwen/Qwen2.5-7B-Instruct": [
31 ([3584, 4608], 1),
32 ([3584, 3584], 0),
33 ([3584, 37888], 1),
34 ([18944, 3584], 0),
35 ],
36 "trymirai/SmolLM2-1.7B-Instruct-AWQ": [
37 ([2048, 6144], 1), # qkv_proj
38 ([2048, 2048], 0), # o_proj
39 ([2048, 16384], 1), # gate_up_proj
40 ([8192, 2048], 0), # down_proj
41 ],
42 "RedHatAI/Qwen3-8B-quantized.w4a16": [
43 ([4096, 6144], 1), # qkv_proj
44 ([4096, 4096], 0), # o_proj
45 ([4096, 24576], 1), # gate_up_proj
46 ([12288, 4096], 0), # down_proj
47 ],
48 }
49
50
51 # ---------------------------------------------------------------------------
52 # Weight packing
53 # ---------------------------------------------------------------------------
54 def prepare_hybrid_weights(K, N, group_size, device="cuda"):
55 """Create random weights for benchmarking.
56
57 Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
58 its int32 view from w_q_skinny, so no separate int32 buffer is returned.
59 """
60 num_groups = K // group_size
61
62 # Random packed weights — actual values don't matter for throughput
63 w_q_skinny_i32 = torch.randint(
64 0, 2**31, (N, K // 8), dtype=torch.int32, device=device
65 )
66 w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
67 w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
68
69 # Raw per-group zero-points for asymmetric benchmarks
70 w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
71 torch.float16
72 )
73
74 # FP16 baseline for F.linear
75 w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
76
77 return w_q_skinny, w_s_skinny, w_fp16, w_zp
78
79
80 # ---------------------------------------------------------------------------
81 # Benchmark
82 # ---------------------------------------------------------------------------
83 PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
84
85
86 @triton.testing.perf_report(
87 triton.testing.Benchmark(
88 x_names=["batch_size"],
89 x_vals=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
90 x_log=False,
91 line_arg="provider",
92 line_vals=PROVIDERS,
93 line_names=PROVIDERS,
94 ylabel="TFLOP/s (larger is better)",
95 plot_name="FP16 vs Hybrid W4A16",
96 args={},
97 )
98 )
99 def benchmark(batch_size, provider, N, K, group_size, weights):
100 M = batch_size
101 device = "cuda"
102 dtype = torch.float16
103 a = torch.randn((M, K), device=device, dtype=dtype)
104
105 quantiles = [0.5, 0.2, 0.8]
106
107 if provider == "torch-fp16":
108 w_fp16 = weights["w_fp16"]
109 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
110 lambda: torch.nn.functional.linear(a, w_fp16),
111 quantiles=quantiles,
112 )
113 elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"):
114 from vllm.model_executor.kernels.linear.mixed_precision import (
115 rdna_hybrid_w4a16 as _k,
116 )
117
118 _rdna_hybrid_w4a16_apply_impl = _k._rdna_hybrid_w4a16_apply_impl
119 from vllm.utils.platform_utils import num_compute_units
120
121 w = weights
122 cu_count = num_compute_units()
123 use_zp = provider == "hybrid-w4a16-zp"
124
125 def run():
126 return _rdna_hybrid_w4a16_apply_impl(
127 a,
128 w["w_q_skinny"],
129 w["w_s_skinny"],
130 w["w_zp"] if use_zp else None,
131 None, # bias
132 cu_count,
133 group_size,
134 )
135
136 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
137 run,
138 quantiles=quantiles,
139 )
140 else:
141 return 0.0, 0.0, 0.0
142
143 to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
144 return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
145
146
147 def prepare_shapes(args):
148 KN_model_names = []
149 for model, tp_size in itertools.product(args.models, args.tp_sizes):
150 for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
151 KN[tp_dim] //= tp_size
152 KN.append(model)
153 KN_model_names.append(KN)
154 return KN_model_names
155
156
157 if __name__ == "__main__":
158 parser = argparse.ArgumentParser(
159 description="Benchmark RDNAHybridW4A16LinearKernel"
160 )
161 parser.add_argument(
162 "--models",
163 nargs="+",
164 type=str,
165 default=["Qwen/Qwen3-4B"],
166 choices=list(WEIGHT_SHAPES.keys()),
167 )
168 parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
169 parser.add_argument("--group-size", type=int, default=128)
170 parser.add_argument("--save-path", type=str, default=None)
171 args = parser.parse_args()
172
173 for K, N, model in prepare_shapes(args):
174 group_size = args.group_size
175 print(f"\n{'=' * 70}")
176 print(f"{model}, N={N} K={K}, group_size={group_size}")
177 print(f"{'=' * 70}")
178
179 w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
180
181 weights = {
182 "w_q_skinny": w_q_skinny,
183 "w_s_skinny": w_s_skinny,
184 "w_fp16": w_fp16,
185 "w_zp": w_zp,
186 }
187
188 save_path = args.save_path or f"bench_int4_res_n{N}_k{K}"
189 os.makedirs(save_path, exist_ok=True)
190 benchmark.run(
191 print_data=True,
192 show_plots=False,
193 save_path=save_path,
194 N=N,
195 K=K,
196 group_size=group_size,
197 weights=weights,
198 )
199
200 print("\nBenchmark finished!")