diff --git a/benchmarks/cutlass_benchmarks/utils.py b/benchmarks/cutlass_benchmarks/utils.py deleted file mode 100644 index 659c68bb11..0000000000 --- a/benchmarks/cutlass_benchmarks/utils.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Cutlass bench utils - -import torch - - -def to_fp8(tensor: torch.Tensor) -> torch.Tensor: - finfo = torch.finfo(torch.float8_e4m3fn) - return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to( - dtype=torch.float8_e4m3fn - ) - - -def to_int8(tensor: torch.Tensor) -> torch.Tensor: - return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8) - - -def to_bf16(tensor: torch.Tensor) -> torch.Tensor: - return tensor.to(dtype=torch.bfloat16) - - -def to_fp16(tensor: torch.Tensor) -> torch.Tensor: - return tensor.to(dtype=torch.float16) - - -def make_rand_tensors( - dtype: torch.dtype, m: int, n: int, k: int -) -> tuple[torch.Tensor, torch.Tensor]: - a = torch.randn((m, k), device="cuda") * 5 - b = torch.randn((n, k), device="cuda").t() * 5 - - if dtype == torch.int8: - return to_int8(a), to_int8(b) - if dtype == torch.float8_e4m3fn: - return to_fp8(a), to_fp8(b) - - raise ValueError("unsupported dtype") diff --git a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py deleted file mode 100644 index f7325ddd2c..0000000000 --- a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py +++ /dev/null @@ -1,372 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import copy -import itertools -import pickle as pkl -import time -from collections.abc import Callable, Iterable - -import torch -import torch.utils.benchmark as TBenchmark -from torch.utils.benchmark import Measurement as TMeasurement -from utils import make_rand_tensors -from weight_shapes import WEIGHT_SHAPES - -from vllm import _custom_ops as ops -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - w8a8_triton_block_scaled_mm, -) -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.utils.math_utils import cdiv - -DEFAULT_MODELS = list(WEIGHT_SHAPES.keys()) -DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512] -DEFAULT_TP_SIZES = [1] - - -# bench -def bench_fn( - label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs -) -> TMeasurement: - min_run_time = 1 - - globals = { - "args": args, - "kwargs": kwargs, - "fn": fn, - } - return TBenchmark.Timer( - stmt="fn(*args, **kwargs)", - globals=globals, - label=label, - sub_label=sub_label, - description=description, - ).blocked_autorange(min_run_time=min_run_time) - - -def bench_int8( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - """Benchmark INT8-based kernels.""" - assert dtype == torch.int8 - a, b = make_rand_tensors(torch.int8, m, n, k) - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32) - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32) - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16) - azp = torch.zeros((m,), device="cuda", dtype=torch.int32) - azp_adj = torch.zeros((n,), device="cuda", dtype=torch.int32) - - bench_fns = { - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16) - ), - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.float16), b.to(dtype=torch.float16) - ), - "cutlass_i8_i8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16 - ), - "cutlass_i8_i8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16, bias - ), - "cutlass_i8_i8_bf16_scaled_mm_azp": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_bias": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, None, bias - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_pt": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp - ), - "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias": lambda: ops.cutlass_scaled_mm_azp( - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp, bias - ), - } - - timers = [] - for name, fn in bench_fns.items(): - # If bench_kernels is None, run all. Otherwise, run only exact matches. - if bench_kernels is None or name in bench_kernels: - print(f"Running {name}") - timers.append(bench_fn(label, sub_label, name, fn)) - - return timers - - -def bench_fp8( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - """Benchmark FP8-based kernels.""" - assert dtype == torch.float8_e4m3fn - a, b = make_rand_tensors(torch.float8_e4m3fn, m, n, k) - a_cont = a.contiguous() - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32) - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32) - - block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32) - block_scale_b = torch.rand( - cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32 - ) - block_scale_a_M_major = block_scale_a.t().contiguous().t() - block_scale_b_K_major = block_scale_b.t().contiguous().t() - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16) - - print(m, k, n) - - bench_fns = { - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16) - ), - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm( - a.to(dtype=torch.float16), b.to(dtype=torch.float16) - ), - "pytorch_fp8_fp8_fp16_scaled_mm": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.float16 - ), - "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.float16, use_fast_accum=True - ), - "pytorch_fp8_fp8_bf16_scaled_mm": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.bfloat16 - ), - "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum": lambda: torch._scaled_mm( - a, b, scale_a, scale_b, out_dtype=torch.bfloat16, use_fast_accum=True - ), - "cutlass_fp8_fp8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16 - ), - "cutlass_fp8_fp8_fp16_scaled_mm": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.float16 - ), - "cutlass_fp8_fp8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.bfloat16, bias - ), - "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm( - a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16) - ), - "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm( - a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128) - ), - "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm( - a, b, block_scale_a_M_major, block_scale_b_K_major, torch.float16 - ), - } - - timers = [] - for name, fn in bench_fns.items(): - # If bench_kernels is None, run all. Otherwise, run only exact matches. - if bench_kernels is None or name in bench_kernels: - print(f"Running {name}") - timers.append(bench_fn(label, sub_label, name, fn)) - - return timers - - -def bench( - dtype: torch.dtype, - m: int, - k: int, - n: int, - label: str, - sub_label: str, - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - if dtype == torch.int8: - return bench_int8(dtype, m, k, n, label, sub_label, bench_kernels) - if dtype == torch.float8_e4m3fn: - return bench_fp8(dtype, m, k, n, label, sub_label, bench_kernels) - raise ValueError("unsupported type") - - -# runner -def print_timers(timers: Iterable[TMeasurement]): - compare = TBenchmark.Compare(timers) - compare.print() - - -def run( - dtype: torch.dtype, - MKNs: Iterable[tuple[int, int, int]], - bench_kernels: list[str] | None = None, -) -> Iterable[TMeasurement]: - results = [] - for m, k, n in MKNs: - timers = bench( - dtype, - m, - k, - n, - f"scaled-{dtype}-gemm", - f"MKN=({m}x{k}x{n})", - bench_kernels=bench_kernels, - ) - print_timers(timers) - results.extend(timers) - return results - - -def make_output( - data: Iterable[TMeasurement], - MKNs: Iterable[tuple[int, int, int]], - base_description: str, - timestamp=None, -): - print(f"== All Results {base_description} ====") - print_timers(data) - - # pickle all the results - timestamp = int(time.time()) if timestamp is None else timestamp - with open(f"{base_description}-{timestamp}.pkl", "wb") as f: - pkl.dump(data, f) - - -def run_square_bench(args): - dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment)) - MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes)) - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - make_output(data, MKNs, f"square_bench-{args.dtype}") - - -def run_range_bench(args): - dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment)) - n = len(dim_sizes) - Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes - Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes - Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes - MKNs = list(zip(Ms, Ks, Ns)) - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - make_output(data, MKNs, f"range_bench-{args.dtype}") - - -def run_model_bench(args): - print("Benchmarking models:") - for i, model in enumerate(args.models): - print(f"[{i}] {model}") - - def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]: - KNs = [] - for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]): - KN[tp_split_dim] = KN[tp_split_dim] // tp_size - KNs.append(KN) - return KNs - - model_bench_data = [] - models_tps = list(itertools.product(args.models, args.tp_sizes)) - for model, tp_size in models_tps: - Ms = args.batch_sizes - KNs = model_shapes(model, tp_size) - MKNs = [] - for m in Ms: - for k, n in KNs: - MKNs.append((m, k, n)) - - data = run(args.dtype, MKNs, bench_kernels=args.kernels) - model_bench_data.append(data) - - # Print all results - for data, model_tp in zip(model_bench_data, models_tps): - model, tp_size = model_tp - print(f"== Results {args.dtype} {model}-TP{tp_size} ====") - print_timers(data) - - timestamp = int(time.time()) - - all_data = [] - for d in model_bench_data: - all_data.extend(d) - # pickle all data - with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f: - pkl.dump(all_data, f) - - -if __name__ == "__main__": - - def to_torch_dtype(dt): - if dt == "int8": - return torch.int8 - if dt == "fp8": - return torch.float8_e4m3fn - raise ValueError("unsupported dtype") - - parser = FlexibleArgumentParser( - description=""" -Benchmark Cutlass GEMM. - - To run square GEMMs: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64 - - To run constant N and K and sweep M: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384 - - To run dimensions from a model: - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1 - - Output: - - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs. - """, # noqa: E501 - formatter_class=argparse.RawTextHelpFormatter, - ) - - parser.add_argument( - "--dtype", - type=to_torch_dtype, - required=True, - help="Available options are ['int8', 'fp8']", - ) - parser.add_argument( - "--kernels", - nargs="+", - type=str, - default=None, - help="Exact names of the kernels to benchmark. If not set, runs all kernels.", - ) - - subparsers = parser.add_subparsers(dest="cmd") - - square_parser = subparsers.add_parser("square_bench") - square_parser.add_argument("--dim-start", type=int, required=True) - square_parser.add_argument("--dim-end", type=int, required=True) - square_parser.add_argument("--dim-increment", type=int, required=True) - square_parser.set_defaults(func=run_square_bench) - - range_parser = subparsers.add_parser("range_bench") - range_parser.add_argument("--dim-start", type=int, required=True) - range_parser.add_argument("--dim-end", type=int, required=True) - range_parser.add_argument("--dim-increment", type=int, required=True) - range_parser.add_argument("--m-constant", type=int, default=None) - range_parser.add_argument("--n-constant", type=int, default=None) - range_parser.add_argument("--k-constant", type=int, default=None) - range_parser.set_defaults(func=run_range_bench) - - model_parser = subparsers.add_parser("model_bench") - model_parser.add_argument( - "--models", - nargs="+", - type=str, - default=DEFAULT_MODELS, - choices=WEIGHT_SHAPES.keys(), - ) - model_parser.add_argument( - "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES - ) - model_parser.add_argument( - "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES - ) - model_parser.set_defaults(func=run_model_bench) - - args = parser.parse_args() - args.func(args) diff --git a/benchmarks/cutlass_benchmarks/weight_shapes.py b/benchmarks/cutlass_benchmarks/weight_shapes.py deleted file mode 100644 index 25b96ef566..0000000000 --- a/benchmarks/cutlass_benchmarks/weight_shapes.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Weight Shapes are in the format -# ([K, N], TP_SPLIT_DIM) -# Example: -# A shape of ([14336, 4096], 0) indicates the following GEMM shape, -# - TP1 : K = 14336, N = 4096 -# - TP2 : K = 7168, N = 4096 -# A shape of ([4096, 6144], 1) indicates the following GEMM shape, -# - TP1 : K = 4096, N = 6144 -# - TP4 : K = 4096, N = 1536 - -# TP1 shapes -WEIGHT_SHAPES = { - "mistralai/Mistral-7B-v0.1": [ - ([4096, 6144], 1), - ([4096, 4096], 0), - ([4096, 28672], 1), - ([14336, 4096], 0), - ], - "meta-llama/Llama-2-7b-hf": [ - ([4096, 12288], 1), - ([4096, 4096], 0), - ([4096, 22016], 1), - ([11008, 4096], 0), - ], - "meta-llama/Llama-3-8b": [ - ([4096, 6144], 1), - ([4096, 4096], 0), - ([4096, 28672], 1), - ([14336, 4096], 0), - ], - "meta-llama/Llama-2-13b-hf": [ - ([5120, 15360], 1), - ([5120, 5120], 0), - ([5120, 27648], 1), - ([13824, 5120], 0), - ], - "meta-llama/Llama-2-70b-hf": [ - ([8192, 10240], 1), - ([8192, 8192], 0), - ([8192, 57344], 1), - ([28672, 8192], 0), - ], -} diff --git a/benchmarks/kernels/benchmark_rmsnorm.py b/benchmarks/kernels/benchmark_rmsnorm.py deleted file mode 100644 index d8d7f5bcf9..0000000000 --- a/benchmarks/kernels/benchmark_rmsnorm.py +++ /dev/null @@ -1,255 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import itertools - -import torch -from flashinfer.norm import fused_add_rmsnorm, rmsnorm -from torch import nn - -from vllm import _custom_ops as vllm_ops -from vllm.triton_utils import triton - - -class HuggingFaceRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - orig_dtype = x.dtype - x = x.to(torch.float32) - if residual is not None: - x = x + residual.to(torch.float32) - residual = x.to(orig_dtype) - - variance = x.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + self.variance_epsilon) - x = x.to(orig_dtype) * self.weight - if residual is None: - return x - else: - return x, residual - - -def rmsnorm_naive( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps) - naive_norm.weight = nn.Parameter(weight) - naive_norm = naive_norm.to(x.device) - - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - output = naive_norm(x, residual) - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def rmsnorm_flashinfer( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - if residual is not None: - fused_add_rmsnorm(x, residual, weight, eps) - output = (x, residual) - else: - output = rmsnorm(x, weight, eps) - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def rmsnorm_vllm( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor | None = None, - eps: float = 1e-6, -): - orig_shape = x.shape - x = x.view(-1, x.shape[-1]) - if residual is not None: - residual = residual.view(-1, residual.shape[-1]) - - if residual is not None: - vllm_ops.fused_add_rms_norm(x, residual, weight, eps) - output = (x, residual) - else: - out = torch.empty_like(x) - vllm_ops.rms_norm(out, x, weight, eps) - output = out - - if isinstance(output, tuple): - output = (output[0].view(orig_shape), output[1].view(orig_shape)) - else: - output = output.view(orig_shape) - return output - - -def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True): - dtype = torch.bfloat16 - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") - weight = torch.ones(hidden_size, dtype=dtype, device="cuda") - residual = torch.randn_like(x) if use_residual else None - - output_naive = rmsnorm_naive( - x.clone(), weight, residual.clone() if residual is not None else None - ) - output_flashinfer = rmsnorm_flashinfer( - x.clone(), weight, residual.clone() if residual is not None else None - ) - output_vllm = rmsnorm_vllm( - x.clone(), weight, residual.clone() if residual is not None else None - ) - - if use_residual: - output_naive = output_naive[0] - output_flashinfer = output_flashinfer[0] - output_vllm = output_vllm[0] - - print(f"Naive output={output_naive}") - print(f"FlashInfer output={output_flashinfer}") - print(f"vLLM output={output_vllm}") - - if torch.allclose( - output_naive, output_flashinfer, atol=1e-2, rtol=1e-2 - ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2): - print("✅ All implementations match") - else: - print("❌ Implementations differ") - - -batch_size_range = [2**i for i in range(0, 7, 2)] -seq_length_range = [2**i for i in range(6, 11, 1)] -head_num_range = [32, 48] -configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range)) - - -def get_benchmark(use_residual): - @triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["head_num", "batch_size", "seq_len"], - x_vals=[list(_) for _ in configs], - line_arg="provider", - line_vals=["huggingface", "flashinfer", "vllm"], - line_names=["HuggingFace", "FlashInfer", "vLLM"], - styles=[("blue", "-"), ("green", "-"), ("red", "-")], - ylabel="us", - plot_name=f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual", - args={}, - ) - ) - def benchmark(head_num, batch_size, seq_len, provider): - dtype = torch.bfloat16 - hidden_size = head_num * 128 # assuming head_dim = 128 - - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") - weight = torch.ones(hidden_size, dtype=dtype, device="cuda") - residual = torch.randn_like(x) if use_residual else None - - quantiles = [0.5, 0.2, 0.8] - - if provider == "huggingface": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_naive( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - elif provider == "flashinfer": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_flashinfer( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - else: - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rmsnorm_vllm( - x.clone(), - weight, - residual.clone() if residual is not None else None, - ), - quantiles=quantiles, - ) - - return 1000 * ms, 1000 * max_ms, 1000 * min_ms - - return benchmark - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument( - "--batch-size", - type=int, - default=4, - help="Batch size", - ) - parser.add_argument( - "--seq-len", - type=int, - default=128, - help="Sequence length", - ) - parser.add_argument( - "--hidden-size", - type=int, - default=4096, - help="Hidden size (2nd dimension) of the sequence", - ) - parser.add_argument( - "--use-residual", action="store_true", help="Whether to use residual connection" - ) - parser.add_argument( - "--save-path", - type=str, - default="./configs/rmsnorm/", - help="Path to save rmsnorm benchmark results", - ) - - args = parser.parse_args() - - # Run correctness test - calculate_diff( - batch_size=args.batch_size, - seq_len=args.seq_len, - hidden_size=args.hidden_size, - use_residual=args.use_residual, - ) - - # Get the benchmark function with proper use_residual setting - benchmark = get_benchmark(args.use_residual) - # Run performance benchmark - benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py index bac2920103..40f5f89476 100644 --- a/benchmarks/kernels/ir/shapes.py +++ b/benchmarks/kernels/ir/shapes.py @@ -18,10 +18,11 @@ COMMON_HIDDEN_SIZES = [ # Each entry maps an op name to a list of kwarg dicts that will be passed # to that op's registered input generator via op.generate_inputs(**kwargs). SHAPE_CONFIGS: dict[str, list[dict]] = { - "rms_norm": [ + op_name: [ {"num_tokens": n, "hidden_size": d, "dtype": dtype} for dtype in [torch.float16, torch.bfloat16, torch.float32] for d in COMMON_HIDDEN_SIZES for n in NUM_TOKENS - ], + ] + for op_name in ("rms_norm", "fused_add_rms_norm") } diff --git a/benchmarks/overheads/benchmark_hashing.py b/benchmarks/overheads/benchmark_hashing.py deleted file mode 100644 index 178599952d..0000000000 --- a/benchmarks/overheads/benchmark_hashing.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import cProfile -import pstats - -from vllm import LLM, SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser - -# A very long prompt, total number of tokens is about 15k. -LONG_PROMPT = ["You are an expert in large language models, aren't you?"] * 1000 -LONG_PROMPT = " ".join(LONG_PROMPT) - - -def main(args): - llm = LLM( - model=args.model, - enforce_eager=True, - enable_prefix_caching=True, - tensor_parallel_size=args.tensor_parallel_size, - ) - - sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len) - profiler = cProfile.Profile() - - print("------warm up------") - for i in range(3): - output = llm.generate(LONG_PROMPT, sampling_params) - print(output[0].outputs[0].text) - - print("------start generating------") - for i in range(3): - profiler.runctx( - "llm.generate(LONG_PROMPT, sampling_params)", globals(), locals() - ) - - # analyze the runtime of hashing function - stats = pstats.Stats(profiler) - stats.sort_stats("cumulative") - total_time = 0 - total_calls = 0 - for func in stats.stats: - if "hash_of_block" in func[2]: - total_time = stats.stats[func][3] - total_calls = stats.stats[func][0] - percentage = (total_time / stats.total_tt) * 100 - print( - f"Hashing took {total_time:.2f} seconds,{percentage:.2f}% of the total runtime." - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser( - description="Benchmark the performance of hashing function in" - "automatic prefix caching." - ) - parser.add_argument("--model", type=str, default="lmsys/longchat-7b-16k") - parser.add_argument("--tensor-parallel-size", "-tp", type=int, default=1) - parser.add_argument("--output-len", type=int, default=10) - parser.add_argument( - "--enable-prefix-caching", action="store_true", help="enable prefix caching" - ) - args = parser.parse_args() - main(args) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 3be7bb243b..f65b9ec921 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -60,8 +60,6 @@ CHECK_IMPORTS = { "benchmarks/kernels/benchmark_lora.py", "benchmarks/kernels/benchmark_machete.py", "benchmarks/fused_kernels/layernorm_rms_benchmarks.py", - "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py", - "benchmarks/cutlass_benchmarks/sparse_benchmarks.py", # cloudpickle "vllm/v1/executor/multiproc_executor.py", "vllm/v1/executor/ray_executor.py",