@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 210 lines 6.26 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
7 import torch
8 from weight_shapes import WEIGHT_SHAPES
9
10 from vllm import _custom_ops as ops
11 from vllm.platforms import current_platform
12 from vllm.scalar_type import scalar_types
13 from vllm.triton_utils import triton
14 from vllm.utils.flashinfer import flashinfer_fp4_quantize
15
16 if not current_platform.has_device_capability(100):
17 raise RuntimeError("NVFP4 requires compute capability of 10.0 (Blackwell)")
18
19 FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
20 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
21
22 PROVIDER_CFGS = {
23 "vllm": dict(backend="vllm", is_sf_swizzled_layout=False, enabled=True),
24 "vllm-swizzle": dict(backend="vllm", is_sf_swizzled_layout=True, enabled=True),
25 "flashinfer": dict(backend="flashinfer", is_sf_swizzled_layout=False, enabled=True),
26 "flashinfer-swizzle": dict(
27 backend="flashinfer", is_sf_swizzled_layout=True, enabled=True
28 ),
29 }
30
31 _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
32
33
34 def compute_global_scale(tensor: torch.Tensor) -> torch.Tensor:
35 """Compute global scale for FP4 quantization."""
36 amax = torch.abs(tensor).max().to(torch.float32)
37 return FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
38
39
40 @triton.testing.perf_report(
41 triton.testing.Benchmark(
42 x_names=["batch_size"],
43 x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
44 x_log=False,
45 line_arg="provider",
46 line_vals=_enabled,
47 line_names=_enabled,
48 ylabel="us (lower is better)",
49 plot_name="NVFP4 Input Quantization Latency (us)",
50 args={},
51 )
52 )
53 def benchmark(batch_size, provider, N, K):
54 M = batch_size
55 device = "cuda"
56 dtype = torch.bfloat16
57
58 # Create input tensor
59 a = torch.randn((M, K), device=device, dtype=dtype)
60
61 # Compute global scale for activation
62 a_global_scale = compute_global_scale(a)
63
64 quantiles = [0.5, 0.2, 0.8]
65
66 cfg = PROVIDER_CFGS[provider]
67
68 if cfg["backend"] == "vllm":
69 # vLLM's FP4 quantization
70 if cfg["is_sf_swizzled_layout"]:
71 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
72 lambda: ops.scaled_fp4_quant(
73 a, a_global_scale, is_sf_swizzled_layout=True
74 ),
75 quantiles=quantiles,
76 )
77 else:
78 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
79 lambda: ops.scaled_fp4_quant(
80 a, a_global_scale, is_sf_swizzled_layout=False
81 ),
82 quantiles=quantiles,
83 )
84 elif cfg["backend"] == "flashinfer":
85 # FlashInfer's FP4 quantization
86 if cfg["is_sf_swizzled_layout"]:
87 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
88 lambda: flashinfer_fp4_quantize(
89 a, a_global_scale, is_sf_swizzled_layout=True
90 ),
91 quantiles=quantiles,
92 )
93 else:
94 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
95 lambda: flashinfer_fp4_quantize(
96 a, a_global_scale, is_sf_swizzled_layout=False
97 ),
98 quantiles=quantiles,
99 )
100
101 # Convert ms to us for better readability at small batch sizes
102 to_us = lambda t_ms: t_ms * 1000
103 return to_us(ms), to_us(max_ms), to_us(min_ms)
104
105
106 def prepare_shapes(args):
107 out = []
108 for model, tp_size in itertools.product(args.models, args.tp_sizes):
109 for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
110 KN[tp_dim] //= tp_size
111 KN.append(model)
112 out.append(KN)
113 return out
114
115
116 def _test_accuracy_once(
117 M: int, K: int, dtype: torch.dtype, device: str, is_sf_swizzled_layout: bool
118 ):
119 """Test accuracy between vLLM and FlashInfer FP4 quantization."""
120 # Create input tensor
121 a = torch.randn((M, K), device=device, dtype=dtype)
122
123 # Compute global scale
124 a_global_scale = compute_global_scale(a)
125
126 # vLLM quantization
127 vllm_fp4, vllm_scale = ops.scaled_fp4_quant(
128 a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
129 )
130
131 # FlashInfer quantization (with swizzled layout to match vLLM's output)
132 flashinfer_fp4, flashinfer_scale = flashinfer_fp4_quantize(
133 a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
134 )
135 flashinfer_scale = flashinfer_scale.view(torch.float8_e4m3fn)
136
137 # Compare outputs
138 torch.testing.assert_close(
139 vllm_fp4,
140 flashinfer_fp4,
141 )
142 # Compare scales
143 torch.testing.assert_close(
144 vllm_scale,
145 flashinfer_scale,
146 )
147 print(
148 f"M={M}, K={K}, dtype={dtype}, is_sf_swizzled_layout={is_sf_swizzled_layout}: PASSED" # noqa: E501
149 )
150
151
152 def test_accuracy():
153 """Run accuracy tests across various shapes."""
154 print("\n" + "=" * 60)
155 print("Running accuracy tests: vLLM vs FlashInfer")
156 print("=" * 60)
157
158 device = "cuda"
159 dtype = torch.bfloat16
160
161 # Test various batch sizes and hidden dimensions
162 Ms = [1, 1024]
163 Ks = [4096]
164
165 for is_sf_swizzled_layout in [True, False]:
166 for M in Ms:
167 for K in Ks:
168 _test_accuracy_once(M, K, dtype, device, is_sf_swizzled_layout)
169
170 print("\nAll accuracy tests passed!")
171
172
173 if __name__ == "__main__":
174 parser = argparse.ArgumentParser(
175 description="Benchmark NVFP4 quantization: vLLM vs FlashInfer"
176 )
177 parser.add_argument(
178 "--models",
179 nargs="+",
180 type=str,
181 default=["meta-llama/Llama-3.3-70B-Instruct"],
182 choices=list(WEIGHT_SHAPES.keys()),
183 )
184 parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
185 parser.add_argument(
186 "--save-path",
187 type=str,
188 default=None,
189 help="Path to save benchmark results",
190 )
191 parser.add_argument(
192 "--accuracy",
193 action="store_true",
194 help="Run accuracy tests",
195 )
196 args = parser.parse_args()
197
198 if args.accuracy:
199 test_accuracy()
200
201 for K, N, model in prepare_shapes(args):
202 print(f"\n{model}, N={N} K={K}")
203 benchmark.run(
204 print_data=True,
205 save_path=args.save_path,
206 N=N,
207 K=K,
208 )
209
210 print("\nBenchmark finished!")