| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | |
| 5 | from benchmark_utils import TimeCollector |
| 6 | from tabulate import tabulate |
| 7 | |
| 8 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 9 | from vllm.v1.core.block_pool import BlockPool |
| 10 | |
| 11 | |
| 12 | def main(args): |
| 13 | rows = [] |
| 14 | for allocate_block in args.allocate_blocks: |
| 15 | # Enforce a GC collect ahead to minimize the impact among runs |
| 16 | gc.collect() |
| 17 | block_pool = BlockPool(num_gpu_blocks=args.num_gpu_blocks, enable_caching=True) |
| 18 | |
| 19 | get_blocks_times = TimeCollector(TimeCollector.US) |
| 20 | free_blocks_times = TimeCollector(TimeCollector.US) |
| 21 | for _ in range(args.num_iteration): |
| 22 | with get_blocks_times: |
| 23 | blocks = block_pool.get_new_blocks(allocate_block) |
| 24 | with free_blocks_times: |
| 25 | block_pool.free_blocks(blocks) |
| 26 | |
| 27 | rows.append( |
| 28 | [get_blocks_times.cnt, args.num_gpu_blocks, allocate_block] |
| 29 | + get_blocks_times.dump_avg_max() |
| 30 | + free_blocks_times.dump_avg_max() |
| 31 | ) |
| 32 | |
| 33 | print( |
| 34 | tabulate( |
| 35 | rows, |
| 36 | headers=[ |
| 37 | "Iterations", |
| 38 | "Total\nBlocks", |
| 39 | "Allocated\nBlocks", |
| 40 | "Get Blocks\nAvg (us)", |
| 41 | "Get Blocks\nMax (us)", |
| 42 | "Free Blocks\nAvg (us)", |
| 43 | "Free Blocks\nMax (us)", |
| 44 | ], |
| 45 | tablefmt="grid", |
| 46 | floatfmt=".3f", |
| 47 | ) |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def invoke_main() -> None: |
| 52 | parser = FlexibleArgumentParser( |
| 53 | description="Benchmark the performance of BlockPool for KV Cache." |
| 54 | ) |
| 55 | parser.add_argument("--num-gpu-blocks", type=int, default=100000) |
| 56 | parser.add_argument( |
| 57 | "--num-iteration", |
| 58 | type=int, |
| 59 | default=1000, |
| 60 | help="Number of iterations to run to stabilize final data readings", |
| 61 | ) |
| 62 | parser.add_argument( |
| 63 | "--allocate-blocks", |
| 64 | type=int, |
| 65 | nargs="*", |
| 66 | default=[10, 50, 100, 500, 1000], |
| 67 | help="Number of blocks to allocate", |
| 68 | ) |
| 69 | args = parser.parse_args() |
| 70 | main(args) |
| 71 | |
| 72 | |
| 73 | if __name__ == "__main__": |
| 74 | invoke_main() # pragma: no cover |