@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 198 lines 6.12 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import argparse
4 import copy
5 import itertools
6 import os
7
8 import torch
9 from weight_shapes import WEIGHT_SHAPES
10
11 from vllm import _custom_ops as ops
12 from vllm.platforms import current_platform
13 from vllm.scalar_type import scalar_types
14 from vllm.triton_utils import triton
15
16 if not current_platform.has_device_capability(100):
17 raise RuntimeError("NVFP4 requires compute capability of 10.0 (Blackwell)")
18
19
20 FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
21 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
22
23 PROVIDER_CFGS = {
24 "torch-bf16": dict(enabled=True),
25 "nvfp4": dict(no_a_quant=False, enabled=True),
26 "nvfp4-noquant": dict(no_a_quant=True, enabled=True),
27 "fbgemm-nvfp4": dict(fbgemm=True, no_a_quant=False, enabled=True),
28 "fbgemm-nvfp4-noquant": dict(fbgemm=True, no_a_quant=True, enabled=True),
29 }
30
31 _needs_fbgemm = any(
32 v.get("fbgemm", False) for v in PROVIDER_CFGS.values() if v.get("enabled", False)
33 )
34 if _needs_fbgemm:
35 try:
36 from fbgemm_gpu.experimental.gemm.triton_gemm.fp4_quantize import (
37 triton_scale_nvfp4_quant,
38 )
39 except ImportError:
40 print(
41 "WARNING: FBGEMM providers are enabled but fbgemm_gpu is not installed. "
42 "These providers will be skipped. Please install fbgemm_gpu with: "
43 "'pip install fbgemm-gpu-genai' to run them."
44 )
45 # Disable FBGEMM providers so the benchmark can run.
46 for cfg in PROVIDER_CFGS.values():
47 if cfg.get("fbgemm"):
48 cfg["enabled"] = False
49
50 _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
51
52
53 def _quant_weight_nvfp4(b: torch.Tensor, device: str, cfg):
54 # Compute global scale for weight
55 b_amax = torch.abs(b).max().to(torch.float32)
56 b_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
57 if "fbgemm" in cfg and cfg["fbgemm"]:
58 b_fp4, scale_b_fp4 = triton_scale_nvfp4_quant(b, b_global_scale)
59 else:
60 b_fp4, scale_b_fp4 = ops.scaled_fp4_quant(b, b_global_scale)
61 return b_fp4, scale_b_fp4, b_global_scale
62
63
64 def build_nvfp4_runner(cfg, a, b, dtype, device):
65 b_fp4, scale_b_fp4, b_global_scale = _quant_weight_nvfp4(b, device, cfg)
66
67 # Compute global scale for activation
68 # NOTE: This is generally provided ahead-of-time by the model checkpoint.
69 a_amax = torch.abs(a).max().to(torch.float32)
70 a_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax
71
72 # Alpha for the GEMM operation
73 alpha = 1.0 / (a_global_scale * b_global_scale)
74 if "fbgemm" in cfg and cfg["fbgemm"]:
75 if cfg["no_a_quant"]:
76 a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
77
78 def run():
79 return torch.ops.fbgemm.f4f4bf16(
80 a_fp4,
81 b_fp4,
82 scale_a_fp4,
83 scale_b_fp4,
84 global_scale=alpha,
85 use_mx=False,
86 )
87
88 return run
89 else:
90
91 def run():
92 a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
93 return torch.ops.fbgemm.f4f4bf16(
94 a_fp4,
95 b_fp4,
96 scale_a_fp4,
97 scale_b_fp4,
98 global_scale=alpha,
99 use_mx=False,
100 )
101
102 return run
103
104 if cfg["no_a_quant"]:
105 # Pre-quantize activation
106 a_fp4, scale_a_fp4 = ops.scaled_fp4_quant(a, a_global_scale)
107
108 def run():
109 return ops.cutlass_scaled_fp4_mm(
110 a_fp4, b_fp4, scale_a_fp4, scale_b_fp4, alpha, dtype
111 )
112
113 return run
114
115 # Quantize activation on-the-fly
116 def run():
117 a_fp4, scale_a_fp4 = ops.scaled_fp4_quant(a, a_global_scale)
118 return ops.cutlass_scaled_fp4_mm(
119 a_fp4, b_fp4, scale_a_fp4, scale_b_fp4, alpha, dtype
120 )
121
122 return run
123
124
125 @triton.testing.perf_report(
126 triton.testing.Benchmark(
127 x_names=["batch_size"],
128 x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
129 x_log=False,
130 line_arg="provider",
131 line_vals=_enabled,
132 line_names=_enabled,
133 ylabel="TFLOP/s (larger is better)",
134 plot_name="BF16 vs NVFP4 GEMMs",
135 args={},
136 )
137 )
138 def benchmark(batch_size, provider, N, K):
139 M = batch_size
140 device = "cuda"
141 dtype = torch.bfloat16
142
143 a = torch.randn((M, K), device=device, dtype=dtype)
144 b = torch.randn((N, K), device=device, dtype=dtype)
145
146 quantiles = [0.5, 0.2, 0.8]
147
148 if provider == "torch-bf16":
149 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
150 lambda: torch.nn.functional.linear(a, b), quantiles=quantiles
151 )
152 else:
153 cfg = PROVIDER_CFGS[provider]
154 run_quant = build_nvfp4_runner(cfg, a, b, dtype, device)
155 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
156 lambda: run_quant(), quantiles=quantiles
157 )
158
159 to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
160 return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
161
162
163 def prepare_shapes(args):
164 out = []
165 for model, tp_size in itertools.product(args.models, args.tp_sizes):
166 for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
167 KN[tp_dim] //= tp_size
168 KN.append(model)
169 out.append(KN)
170 return out
171
172
173 if __name__ == "__main__":
174 parser = argparse.ArgumentParser()
175 parser.add_argument(
176 "--models",
177 nargs="+",
178 type=str,
179 default=["meta-llama/Llama-3.1-8B-Instruct"],
180 choices=list(WEIGHT_SHAPES.keys()),
181 )
182 parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
183 args = parser.parse_args()
184
185 for K, N, model in prepare_shapes(args):
186 print(f"{model}, N={N} K={K}, BF16 vs NVFP4 GEMMs TFLOP/s:")
187 save_dir = f"bench_nvfp4_res_n{N}_k{K}"
188 os.makedirs(save_dir, exist_ok=True)
189
190 benchmark.run(
191 print_data=True,
192 show_plots=True,
193 save_path=save_dir,
194 N=N,
195 K=K,
196 )
197
198 print("Benchmark finished!")