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",