mirror of https://github.com/vllm-project/vllm.git
65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
# 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 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)
|