| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Offline benchmark to test the long document QA throughput. |
| 4 | |
| 5 | Example usage: |
| 6 | # This workload samples 8 different prompts with a default input |
| 7 | # length of 20000 tokens, then replicates each prompt 2 times |
| 8 | # in random order. |
| 9 | python benchmark_long_document_qa_throughput.py \ |
| 10 | --model meta-llama/Llama-2-7b-chat-hf \ |
| 11 | --enable-prefix-caching \ |
| 12 | --num-documents 8 \ |
| 13 | --repeat-count 2 |
| 14 | |
| 15 | Commandline arguments: |
| 16 | --num-documents: The number of documents to sample prompts from. |
| 17 | |
| 18 | --document-length: The length of each document in tokens. |
| 19 | (Optional, default: 20000) |
| 20 | |
| 21 | --output-len: The number of tokens to generate for each prompt. |
| 22 | (Optional, default: 10) |
| 23 | |
| 24 | --repeat-count: The number of times to repeat each prompt. |
| 25 | (Optional, default: 2) |
| 26 | |
| 27 | --repeat-mode: The mode to repeat prompts. The supported modes are: |
| 28 | - 'random': shuffle the prompts randomly. (Default) |
| 29 | - 'tile': the entire prompt list is repeated in sequence. (Potentially |
| 30 | lowest cache hit) |
| 31 | - 'interleave': each prompt is repeated consecutively before |
| 32 | moving to the next element. (Highest cache hit) |
| 33 | |
| 34 | --shuffle-seed: Random seed when the repeat mode is "random". |
| 35 | (Optional, default: 0) |
| 36 | |
| 37 | In the meantime, it also supports all the vLLM engine args to initialize the |
| 38 | LLM engine. You can refer to the `vllm.engine.arg_utils.EngineArgs` for more |
| 39 | details. |
| 40 | """ |
| 41 | |
| 42 | import random |
| 43 | import time |
| 44 | |
| 45 | from vllm import LLM, SamplingParams |
| 46 | from vllm.engine.arg_utils import EngineArgs |
| 47 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 48 | |
| 49 | |
| 50 | def test_long_document_qa(llm=None, sampling_params=None, prompts=None): |
| 51 | """Test long document QA with the given prompts and sampling parameters. |
| 52 | Print the time spent in processing all the prompts. |
| 53 | |
| 54 | Args: |
| 55 | llm: The language model used for generating responses. |
| 56 | sampling_params: Sampling parameter used to generate the response. |
| 57 | prompts: A list of prompt strings to be processed by the LLM. |
| 58 | |
| 59 | """ |
| 60 | start_time = time.time() |
| 61 | llm.generate(prompts, sampling_params=sampling_params) |
| 62 | end_time = time.time() |
| 63 | print(f"Time to execute all requests: {end_time - start_time:.4f} secs") |
| 64 | |
| 65 | |
| 66 | def repeat_prompts(prompts, repeat_count, mode: str): |
| 67 | """Repeat each prompt in the list for a specified number of times. |
| 68 | The order of prompts in the output list depends on the mode. |
| 69 | |
| 70 | Args: |
| 71 | prompts: A list of prompts to be repeated. |
| 72 | repeat_count: The number of times each prompt is repeated. |
| 73 | mode: The mode of repetition. Supported modes are: |
| 74 | - 'random': Shuffle the prompts randomly after repetition. |
| 75 | - 'tile': Repeat the entire prompt list in sequence. |
| 76 | Example: [1, 2, 3] -> [1, 2, 3, 1, 2, 3]. |
| 77 | - 'interleave': Repeat each prompt consecutively before moving to |
| 78 | the next. Example: [1, 2, 3] -> [1, 1, 2, 2, 3, 3]. |
| 79 | |
| 80 | Returns: |
| 81 | A list of repeated prompts in the specified order. |
| 82 | |
| 83 | Raises: |
| 84 | ValueError: If an invalid mode is provided. |
| 85 | |
| 86 | """ |
| 87 | print("Repeat mode: ", mode) |
| 88 | if mode == "random": |
| 89 | repeated_prompts = prompts * repeat_count |
| 90 | random.shuffle(repeated_prompts) |
| 91 | return repeated_prompts |
| 92 | elif mode == "tile": |
| 93 | return prompts * repeat_count |
| 94 | elif mode == "interleave": |
| 95 | repeated_prompts = [] |
| 96 | for prompt in prompts: |
| 97 | repeated_prompts.extend([prompt] * repeat_count) |
| 98 | return repeated_prompts |
| 99 | else: |
| 100 | raise ValueError( |
| 101 | f"Invalid mode: {mode}, only support 'random', 'tile', 'interleave'" |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | def main(args): |
| 106 | random.seed(args.shuffle_seed) |
| 107 | |
| 108 | # Prepare the prompts: |
| 109 | # we append the document id at the beginning to avoid any of the document |
| 110 | # being the prefix of other documents |
| 111 | prompts = [ |
| 112 | str(i) + " ".join(["hi"] * args.document_length) |
| 113 | for i in range(args.num_documents) |
| 114 | ] |
| 115 | |
| 116 | prompts = repeat_prompts(prompts, args.repeat_count, mode=args.repeat_mode) |
| 117 | |
| 118 | warmup_prompts = [ |
| 119 | "This is warm up request " + str(i) + " ".join(["hi"] * args.document_length) |
| 120 | for i in range(args.num_documents) |
| 121 | ] |
| 122 | |
| 123 | # Create the LLM engine |
| 124 | engine_args = EngineArgs.from_cli_args(args) |
| 125 | llm = LLM.from_engine_args(engine_args) |
| 126 | sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len) |
| 127 | |
| 128 | print("------warm up------") |
| 129 | test_long_document_qa( |
| 130 | llm=llm, |
| 131 | prompts=warmup_prompts, |
| 132 | sampling_params=sampling_params, |
| 133 | ) |
| 134 | |
| 135 | print("------start generating------") |
| 136 | test_long_document_qa( |
| 137 | llm=llm, |
| 138 | prompts=prompts, |
| 139 | sampling_params=sampling_params, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | def create_argument_parser(): |
| 144 | parser = FlexibleArgumentParser( |
| 145 | description="Benchmark the performance with or " |
| 146 | "without automatic prefix caching." |
| 147 | ) |
| 148 | |
| 149 | parser.add_argument( |
| 150 | "--document-length", |
| 151 | type=int, |
| 152 | # Roughly the number of tokens for a system paper, |
| 153 | # excluding images |
| 154 | default=20000, |
| 155 | help="Range of input lengths for sampling prompts, " |
| 156 | 'specified as "min:max" (e.g., "128:256").', |
| 157 | ) |
| 158 | |
| 159 | parser.add_argument( |
| 160 | "--num-documents", |
| 161 | type=int, |
| 162 | default=8, |
| 163 | help="Range of input lengths for sampling prompts, " |
| 164 | 'specified as "min:max" (e.g., "128:256").', |
| 165 | ) |
| 166 | |
| 167 | parser.add_argument("--output-len", type=int, default=10) |
| 168 | |
| 169 | parser.add_argument( |
| 170 | "--repeat-count", |
| 171 | type=int, |
| 172 | default=2, |
| 173 | help="Number of times to repeat each prompt", |
| 174 | ) |
| 175 | |
| 176 | parser.add_argument( |
| 177 | "--repeat-mode", |
| 178 | type=str, |
| 179 | default="random", |
| 180 | help="The mode to repeat prompts. The supported " |
| 181 | 'modes are "random", "tile", and "interleave". ' |
| 182 | "See repeat_prompts() in the source code for details.", |
| 183 | ) |
| 184 | |
| 185 | parser.add_argument( |
| 186 | "--shuffle-seed", |
| 187 | type=int, |
| 188 | default=0, |
| 189 | help='Random seed when the repeat mode is "random"', |
| 190 | ) |
| 191 | |
| 192 | parser = EngineArgs.add_cli_args(parser) |
| 193 | |
| 194 | return parser |
| 195 | |
| 196 | |
| 197 | if __name__ == "__main__": |
| 198 | parser = create_argument_parser() |
| 199 | args = parser.parse_args() |
| 200 | main(args) |