A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Fixes and updates to bench_per_token_quant_fp8 (#25591)
Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Sep 24, 2025 at 11:30 UTC
d83f3f7cb37a0f1861f16c84d529abcd54889885
1 file changed
+26
-24
benchmarks/kernels/bench_per_token_quant_fp8.py
+26
-24
@@ -51,7 +51,7 @@ def calculate_diff(
51
):
52
"""Calculate the difference between Inductor and CUDA implementations."""
53
device = torch.device("cuda")
54
- x = torch.rand((batch_size * hidden_size, 4096), dtype=dtype, device=device)
54
+ x = torch.randn((batch_size, hidden_size), dtype=dtype, device=device)
55
56
quant_fp8 = QuantFP8(False, group_shape, column_major_scales=False)
57
@@ -59,23 +59,25 @@ def calculate_diff(
59
torch_eager_out, torch_eager_scale = quant_fp8.forward_native(x)
60
cuda_out, cuda_scale = quant_fp8.forward_cuda(x)
61
62
- out_allclose = lambda o1, o2: torch.allclose(
63
- o1.to(torch.float32),
64
- o2.to(torch.float32),
65
- rtol=1e-3,
66
- atol=1e-5,
67
- )
68
- scale_allclose = lambda s1, s2: torch.allclose(s1, s2, rtol=1e-3, atol=1e-5)
69
-
70
- if (
71
- out_allclose(cuda_out, torch_out)
72
- and scale_allclose(cuda_scale, torch_scale)
73
- and out_allclose(cuda_out, torch_eager_out)
74
- and scale_allclose(cuda_scale, torch_eager_scale)
75
- ):
62
+ try:
63
+ torch.testing.assert_close(
64
+ cuda_out.to(torch.float32),
65
+ torch_out.to(torch.float32),
66
+ rtol=1e-3,
67
+ atol=1e-5,
68
+ )
69
+ torch.testing.assert_close(cuda_scale, torch_scale, rtol=1e-3, atol=1e-5)
70
+ torch.testing.assert_close(
71
+ cuda_out.to(torch.float32),
72
+ torch_eager_out.to(torch.float32),
73
+ rtol=1e-3,
74
+ atol=1e-5,
75
+ )
76
+ torch.testing.assert_close(cuda_scale, torch_eager_scale, rtol=1e-3, atol=1e-5)
77
print("✅ All implementations match")
77
- else:
78
+ except AssertionError as e:
79
print("❌ Implementations differ")
80
+ print(e)
81
82
83
configs = []
@@ -91,7 +93,7 @@ def benchmark_quantization(
93
):
94
device = torch.device("cuda")
95
94
- x = torch.randn(batch_size * hidden_size, 4096, device=device, dtype=dtype)
96
+ x = torch.randn(batch_size, hidden_size, device=device, dtype=dtype)
97
98
quantiles = [0.5, 0.2, 0.8]
99
quant_fp8 = QuantFP8(False, group_shape, column_major_scales=col_major)
@@ -157,21 +159,21 @@ if __name__ == "__main__":
159
)
160
parser.add_argument("-c", "--check", action="store_true")
161
parser.add_argument(
160
- "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"
162
+ "--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16"
163
)
164
parser.add_argument(
165
"--hidden-sizes",
166
type=int,
167
nargs="+",
166
- default=None,
167
- help="Hidden sizes to benchmark (default: 1,16,64,128,256,512,1024,2048,4096)",
168
+ default=[896, 1024, 2048, 4096, 7168],
169
+ help="Hidden sizes to benchmark",
170
)
171
parser.add_argument(
172
"--batch-sizes",
173
type=int,
174
nargs="+",
173
- default=None,
174
- help="Batch sizes to benchmark (default: 1,16,32,64,128)",
175
+ default=[1, 16, 128, 512, 1024],
176
+ help="Batch sizes to benchmark",
177
)
178
parser.add_argument(
179
"--group-sizes",
@@ -192,8 +194,8 @@ if __name__ == "__main__":
194
195
dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
196
195
- hidden_sizes = args.hidden_sizes or [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
196
- batch_sizes = args.batch_sizes or [1, 16, 32, 64, 128]
197
+ hidden_sizes = args.hidden_sizes
198
+ batch_sizes = args.batch_sizes
199
200
if args.group_sizes is not None:
201
group_shapes = []