@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 23114d3364
Fork 0

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 15, 17:15 UTC

[Misc] Warn about v0 in benchmark_paged_attn.py (#15495)

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Tyler Michael Smith committed Mar 25, 2025 at 23:31 UTC 23114d33640175229a395b9ed1128c3a41ad65d9
1 file changed +6
benchmarks/kernels/benchmark_paged_attention.py
+6
@@ -7,10 +7,13 @@ from typing import Optional
7 import torch
8
9 from vllm import _custom_ops as ops
10 +from vllm.logger import init_logger
11 from vllm.platforms import current_platform
12 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser,
13 create_kv_caches_with_random)
14
15 +logger = init_logger(__name__)
16 +
17 NUM_BLOCKS = 128 * 1024
18 PARTITION_SIZE = 512
19 PARTITION_SIZE_ROCM = 256
@@ -193,6 +196,9 @@ def main(
196
197
198 if __name__ == '__main__':
199 + logger.warning("This script benchmarks the paged attention kernel. "
200 + "By default this is no longer used in vLLM inference.")
201 +
202 parser = FlexibleArgumentParser(
203 description="Benchmark the paged attention kernel.")
204 parser.add_argument("--version",