95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
import argparse
|
|
|
|
import torch
|
|
import triton
|
|
from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant
|
|
|
|
from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
|
|
|
|
|
|
@torch.compile(backend="inductor")
|
|
def torch_int8_quant(x):
|
|
int8_max = torch.iinfo(torch.int8).max
|
|
|
|
abs_max = x.abs().max(dim=-1, keepdim=True).values
|
|
scales = abs_max.to(torch.float32) / float(int8_max)
|
|
|
|
q_x = (x / scales).round().to(torch.int8)
|
|
|
|
return q_x, scales
|
|
|
|
|
|
def _test_accuracy_once(M, K, input_dtype, device):
|
|
x = torch.randn(M, K, dtype=input_dtype, device=device) * 5000
|
|
out, scales, _ = vllm_scaled_int8_quant(x, symmetric=True)
|
|
out1, scales1 = per_token_quant_int8(x)
|
|
out2, scales2 = torch_int8_quant(x)
|
|
torch.testing.assert_close(out, out2, atol=1, rtol=0)
|
|
torch.testing.assert_close(out, out1, atol=1, rtol=0)
|
|
torch.testing.assert_close(scales, scales2)
|
|
torch.testing.assert_close(scales1, scales2)
|
|
print(f"M: {M}, K: {K}, type: {input_dtype} OK")
|
|
|
|
|
|
def test_accuracy():
|
|
Ms = [1, 13, 128, 1024, 2048, 4096]
|
|
Ks = [512, 1024, 2048, 8192]
|
|
input_dtypes = [torch.float16, torch.bfloat16]
|
|
for M in Ms:
|
|
for K in Ks:
|
|
for input_dtype in input_dtypes:
|
|
_test_accuracy_once(M, K, input_dtype, "cuda")
|
|
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=["batch_size"],
|
|
x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048],
|
|
x_log=False,
|
|
line_arg="provider",
|
|
line_vals=["vllm op", "triton", "torch.compile"],
|
|
line_names=["vllm op", "triton", "torch.compile"],
|
|
styles=[("blue", "-"), ("orange", "-"), ("red", "-")],
|
|
ylabel="ms",
|
|
plot_name="int8 per token quant",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark(batch_size, provider):
|
|
M, K = batch_size, 16384
|
|
x = torch.randn(M, K, dtype=torch.float16, device="cuda") * 1000
|
|
|
|
quantiles = [0.5, 0.2, 0.8]
|
|
if provider == "vllm op":
|
|
ms, min_ms, max_ms = triton.testing.do_bench(
|
|
lambda: vllm_scaled_int8_quant(x, symmetric=True),
|
|
quantiles=quantiles,
|
|
)
|
|
if provider == "triton":
|
|
ms, min_ms, max_ms = triton.testing.do_bench(
|
|
lambda: per_token_quant_int8(x),
|
|
quantiles=quantiles,
|
|
)
|
|
if provider == "torch.compile":
|
|
ms, min_ms, max_ms = triton.testing.do_bench(
|
|
lambda: torch_int8_quant(x),
|
|
quantiles=quantiles,
|
|
)
|
|
|
|
return ms, min_ms, max_ms
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--save_path",
|
|
type=str,
|
|
default="./bench_int8_quant_res",
|
|
help="Path to save int8 quant benchmark results",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
test_accuracy()
|
|
|
|
benchmark.run(print_data=True, show_plots=True, save_path=args.save_path)
|