| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Micro benchmark comparing built-in hash(), SHA-256, and xxHash. |
| 4 | |
| 5 | This focuses on a single test payload shaped like the prefix-cache hash input: |
| 6 | (32-byte bytes object, 32-int tuple) |
| 7 | |
| 8 | Usage: |
| 9 | python benchmarks/hash_micro_benchmark.py --iterations 20000 |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import argparse |
| 15 | import random |
| 16 | import statistics |
| 17 | import time |
| 18 | from collections.abc import Callable, Iterable |
| 19 | |
| 20 | from vllm.utils.hashing import sha256, xxhash |
| 21 | |
| 22 | |
| 23 | def _generate_test_data(seed: int) -> tuple[bytes, tuple[int, ...]]: |
| 24 | """Generate a deterministic test payload.""" |
| 25 | random.seed(seed) |
| 26 | bytes_data = bytes(random.getrandbits(8) for _ in range(32)) |
| 27 | int_tuple = tuple(random.randint(1, 1_000_000) for _ in range(32)) |
| 28 | return (bytes_data, int_tuple) |
| 29 | |
| 30 | |
| 31 | def _benchmark_func(func: Callable[[tuple], object], data: tuple, iterations: int): |
| 32 | """Return (avg_seconds, std_seconds) for hashing `data` `iterations` times.""" |
| 33 | times: list[float] = [] |
| 34 | |
| 35 | # Warm-up to avoid first-run noise. |
| 36 | for _ in range(200): |
| 37 | func(data) |
| 38 | |
| 39 | for _ in range(iterations): |
| 40 | start = time.perf_counter() |
| 41 | func(data) |
| 42 | end = time.perf_counter() |
| 43 | times.append(end - start) |
| 44 | |
| 45 | avg = statistics.mean(times) |
| 46 | std = statistics.stdev(times) if len(times) > 1 else 0.0 |
| 47 | return avg, std |
| 48 | |
| 49 | |
| 50 | def _run_benchmarks( |
| 51 | benchmarks: Iterable[tuple[str, Callable[[tuple], object]]], |
| 52 | data: tuple, |
| 53 | iterations: int, |
| 54 | ): |
| 55 | """Yield (name, avg, std) for each benchmark, skipping unavailable ones.""" |
| 56 | for name, func in benchmarks: |
| 57 | try: |
| 58 | avg, std = _benchmark_func(func, data, iterations) |
| 59 | except ModuleNotFoundError as exc: |
| 60 | print(f"Skipping {name}: {exc}") |
| 61 | continue |
| 62 | yield name, avg, std |
| 63 | |
| 64 | |
| 65 | def builtin_hash(data: tuple) -> int: |
| 66 | """Wrapper for Python's built-in hash().""" |
| 67 | return hash(data) |
| 68 | |
| 69 | |
| 70 | def main() -> None: |
| 71 | parser = argparse.ArgumentParser(description=__doc__) |
| 72 | parser.add_argument( |
| 73 | "--iterations", |
| 74 | type=int, |
| 75 | default=10_000, |
| 76 | help="Number of measured iterations per hash function.", |
| 77 | ) |
| 78 | parser.add_argument( |
| 79 | "--seed", type=int, default=42, help="Random seed for test payload." |
| 80 | ) |
| 81 | args = parser.parse_args() |
| 82 | |
| 83 | data = _generate_test_data(args.seed) |
| 84 | benchmarks = ( |
| 85 | ("SHA256 (pickle)", sha256), |
| 86 | ("xxHash (pickle)", xxhash), |
| 87 | ("built-in hash()", builtin_hash), |
| 88 | ) |
| 89 | |
| 90 | print("=" * 60) |
| 91 | print("HASH FUNCTION MICRO BENCHMARK") |
| 92 | print("=" * 60) |
| 93 | print("Test data: (32-byte bytes object, 32-int tuple)") |
| 94 | print(f"Iterations: {args.iterations:,}") |
| 95 | print("=" * 60) |
| 96 | |
| 97 | results = list(_run_benchmarks(benchmarks, data, args.iterations)) |
| 98 | builtin_entry = next((r for r in results if r[0] == "built-in hash()"), None) |
| 99 | |
| 100 | print("\nResults:") |
| 101 | for name, avg, std in results: |
| 102 | print(f" {name:16s}: {avg * 1e6:8.2f} ± {std * 1e6:6.2f} μs") |
| 103 | |
| 104 | if builtin_entry: |
| 105 | _, builtin_avg, _ = builtin_entry |
| 106 | print("\n" + "=" * 60) |
| 107 | print("SUMMARY (relative to built-in hash())") |
| 108 | print("=" * 60) |
| 109 | for name, avg, _ in results: |
| 110 | if name == "built-in hash()": |
| 111 | continue |
| 112 | speed_ratio = avg / builtin_avg |
| 113 | print(f"• {name} is {speed_ratio:.1f}x slower than built-in hash()") |
| 114 | else: |
| 115 | print("\nBuilt-in hash() result missing; cannot compute speed ratios.") |
| 116 | |
| 117 | |
| 118 | if __name__ == "__main__": |
| 119 | main() |