@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 207 lines 6.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 #
4 # Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 # All Rights Reserved.
6 #
7 # Licensed under the Apache License, Version 2.0 (the "License");
8 # you may not use this file except in compliance with the License.
9 # You may obtain a copy of the License at
10 #
11 # http://www.apache.org/licenses/LICENSE-2.0
12 #
13 # Unless required by applicable law or agreed to in writing, software
14 # distributed under the License is distributed on an "AS IS" BASIS,
15 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 # See the License for the specific language governing permissions and
17 # limitations under the License.
18 #
19
20 import argparse
21 import copy
22 import itertools
23
24 import torch
25 from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
26 from weight_shapes import WEIGHT_SHAPES
27
28 from vllm import _custom_ops as ops # use existing nvfp4 gemm in vllm
29 from vllm._custom_ops import fusedQuantizeNv
30 from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
31 from vllm.triton_utils import triton
32
33 PROVIDER_CFGS = {
34 "torch-bf16": dict(enabled=True),
35 "nvfp4": dict(no_a_quant=False, enabled=True),
36 "nvfp4-noquant": dict(no_a_quant=True, enabled=True),
37 }
38
39 _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
40
41
42 def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
43 return (
44 deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
45 * group_size**-0.5
46 )
47
48
49 def _quant_weight_nvfp4(
50 b: torch.Tensor,
51 forward_hadamard_matrix: torch.Tensor,
52 global_scale: torch.Tensor,
53 device: str,
54 M: int,
55 N: int,
56 K: int,
57 ):
58 weight_hf_e2m1, weight_hf_e8m0 = fusedQuantizeNv(
59 b, forward_hadamard_matrix, global_scale
60 )
61 weight_hf_scale_block = to_blocked(weight_hf_e8m0, backend="triton").view(
62 -1, K // 16
63 )
64 return weight_hf_e2m1, weight_hf_scale_block
65
66
67 def build_nvfp4_runner(cfg, a, b, forward_hadamard_matrix, dtype, device, M, N, K):
68 alpha = torch.tensor([1.0], device="cuda")
69 global_scale = torch.tensor([1.0], device="cuda")
70 weight_hf_e2m1, weight_hf_scale_block = _quant_weight_nvfp4(
71 b, forward_hadamard_matrix, global_scale, device, M, N, K
72 )
73
74 if cfg["no_a_quant"]:
75 # Pre-quantize activation
76 input_hf_e2m1, input_hf_e8m0 = fusedQuantizeNv(
77 a, forward_hadamard_matrix, global_scale
78 )
79 input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton").view(
80 -1, K // 16
81 )
82
83 def run():
84 return ops.cutlass_scaled_fp4_mm(
85 input_hf_e2m1,
86 weight_hf_e2m1,
87 input_hf_scale_block,
88 weight_hf_scale_block,
89 alpha,
90 torch.bfloat16,
91 )
92
93 return run
94
95 # Quantize activation on-the-fly
96 def run():
97 input_hf_e2m1, input_hf_e8m0 = fusedQuantizeNv(
98 a, forward_hadamard_matrix, global_scale
99 )
100 input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton").view(
101 -1, K // 16
102 )
103 return ops.cutlass_scaled_fp4_mm(
104 input_hf_e2m1,
105 weight_hf_e2m1,
106 input_hf_scale_block,
107 weight_hf_scale_block,
108 alpha,
109 torch.bfloat16,
110 )
111
112 return run
113
114
115 @triton.testing.perf_report(
116 triton.testing.Benchmark(
117 x_names=["batch_size"],
118 x_vals=[
119 1,
120 4,
121 8,
122 16,
123 32,
124 64,
125 128,
126 256,
127 512,
128 1024,
129 2048,
130 4096,
131 8192,
132 16384,
133 24576,
134 32768,
135 ],
136 x_log=False,
137 line_arg="provider",
138 line_vals=_enabled,
139 line_names=_enabled,
140 ylabel="TFLOP/s (larger is better)",
141 plot_name="BF16 vs NVFP4 GEMMs",
142 args={},
143 )
144 )
145 def benchmark(batch_size, provider, N, K, had_size):
146 M = batch_size
147 device = "cuda"
148 dtype = torch.bfloat16
149
150 a = torch.randn((M, K), device=device, dtype=dtype)
151 b = torch.randn((N, K), device=device, dtype=dtype)
152 forward_hadamard_matrix = get_hadamard_matrix(had_size, dtype, device)
153
154 quantiles = [0.5, 0.2, 0.8]
155
156 if provider == "torch-bf16":
157 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
158 lambda: torch.nn.functional.linear(a, b), rep=200, quantiles=quantiles
159 )
160 else:
161 cfg = PROVIDER_CFGS[provider]
162 run_quant = build_nvfp4_runner(
163 cfg, a, b, forward_hadamard_matrix, dtype, device, M, N, K
164 )
165 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
166 lambda: run_quant(), rep=200, quantiles=quantiles
167 )
168
169 to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
170 return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
171
172
173 def prepare_shapes(args):
174 out = []
175 for model, tp_size in itertools.product(args.models, args.tp_sizes):
176 for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
177 KN[tp_dim] //= tp_size
178 KN.append(model)
179 out.append(KN)
180 return out
181
182
183 if __name__ == "__main__":
184 parser = argparse.ArgumentParser()
185 parser.add_argument(
186 "--models",
187 nargs="+",
188 type=str,
189 default=["meta-llama/Llama-3.3-70B-Instruct"],
190 choices=list(WEIGHT_SHAPES.keys()),
191 )
192 parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
193 args = parser.parse_args()
194
195 for K, N, model in prepare_shapes(args):
196 for had_size in [16, 32, 64, 128]:
197 print(f"{model}, N={N} K={K}, HAD={had_size}, BF16 vs NVFP4 GEMMs TFLOP/s:")
198 benchmark.run(
199 print_data=True,
200 show_plots=True,
201 save_path=f"bench_nvfp4_res_n{N}_k{K}",
202 N=N,
203 K=K,
204 had_size=had_size,
205 )
206
207 print("Benchmark finished!")