throughput.py 28.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Benchmark offline inference throughput."""
4

5
6
7
8
9
10
11
import argparse
import dataclasses
import json
import os
import random
import time
import warnings
12
from typing import Any
13
14
15
16

import torch
import uvloop
from tqdm import tqdm
17
from transformers import AutoModelForCausalLM, PreTrainedTokenizerBase
18
19
20
21
22
23

from vllm.benchmarks.datasets import (
    AIMODataset,
    BurstGPTDataset,
    ConversationDataset,
    InstructCoderDataset,
24
    MultiModalConversationDataset,
25
26
27
28
29
30
31
32
    PrefixRepetitionRandomDataset,
    RandomDataset,
    SampleRequest,
    ShareGPTDataset,
    SonnetDataset,
    VisionArenaDataset,
)
from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json
33
34
35
36
37
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
from vllm.inputs import TextPrompt, TokensPrompt
from vllm.lora.request import LoRARequest
from vllm.outputs import RequestOutput
from vllm.sampling_params import BeamSearchParams
38
from vllm.tokenizers import TokenizerLike, get_tokenizer
39
from vllm.utils.async_utils import merge_async_iterators
40
41
42
43
44
45


def run_vllm(
    requests: list[SampleRequest],
    n: int,
    engine_args: EngineArgs,
46
    do_profile: bool,
47
    disable_detokenize: bool = False,
48
) -> tuple[float, list[RequestOutput] | None]:
49
    from vllm import LLM, SamplingParams
50

51
52
    llm = LLM(**dataclasses.asdict(engine_args))
    assert all(
53
54
55
56
57
58
59
        llm.llm_engine.model_config.max_model_len
        >= (request.prompt_len + request.expected_output_len)
        for request in requests
    ), (
        "Please ensure that max_model_len is greater than the sum of"
        " prompt_len and expected_output_len for all requests."
    )
60
    # Add the requests to the engine.
61
    prompts: list[TextPrompt | TokensPrompt] = []
62
63
    sampling_params: list[SamplingParams] = []
    for request in requests:
64
65
        prompt = (
            TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"])
66
            if "prompt_token_ids" in request.prompt
67
            else TextPrompt(prompt=request.prompt)
68
        )
69
70
71
72
73
        if request.multi_modal_data:
            assert isinstance(request.multi_modal_data, dict)
            prompt["multi_modal_data"] = request.multi_modal_data
        prompts.append(prompt)

74
75
76
77
78
79
80
81
        sampling_params.append(
            SamplingParams(
                n=n,
                temperature=1.0,
                top_p=1.0,
                ignore_eos=True,
                max_tokens=request.expected_output_len,
                detokenize=not disable_detokenize,
82
83
            )
        )
84
    lora_requests: list[LoRARequest] | None = None
85
86
87
88
89
90
91
92
    if engine_args.enable_lora:
        lora_requests = [request.lora_request for request in requests]

    use_beam_search = False

    outputs = None
    if not use_beam_search:
        start = time.perf_counter()
93
94
        if do_profile:
            llm.start_profile()
95
96
97
        outputs = llm.generate(
            prompts, sampling_params, lora_request=lora_requests, use_tqdm=True
        )
98
99
        if do_profile:
            llm.stop_profile()
100
101
102
103
104
        end = time.perf_counter()
    else:
        assert lora_requests is None, "BeamSearch API does not support LoRA"
        prompts = [request.prompt for request in requests]
        # output_len should be the same for all requests.
105
        output_len = requests[0].expected_output_len
106
107
108
        for request in requests:
            assert request.expected_output_len == output_len
        start = time.perf_counter()
109
110
        if do_profile:
            llm.start_profile()
111
112
113
114
115
116
        llm.beam_search(
            prompts,
            BeamSearchParams(
                beam_width=n,
                max_tokens=output_len,
                ignore_eos=True,
117
118
            ),
        )
119
120
        if do_profile:
            llm.stop_profile()
121
122
123
124
125
        end = time.perf_counter()
    return end - start, outputs


def run_vllm_chat(
126
127
128
129
130
131
    requests: list[SampleRequest],
    n: int,
    engine_args: EngineArgs,
    do_profile: bool,
    disable_detokenize: bool = False,
) -> tuple[float, list[RequestOutput]]:
132
133
134
135
136
137
    """
    Run vLLM chat benchmark. This function is recommended ONLY for benchmarking
    multimodal models as it properly handles multimodal inputs and chat
    formatting. For non-multimodal models, use run_vllm() instead.
    """
    from vllm import LLM, SamplingParams
138

139
140
141
    llm = LLM(**dataclasses.asdict(engine_args))

    assert all(
142
143
144
145
146
147
148
        llm.llm_engine.model_config.max_model_len
        >= (request.prompt_len + request.expected_output_len)
        for request in requests
    ), (
        "Please ensure that max_model_len is greater than the sum of "
        "prompt_len and expected_output_len for all requests."
    )
149
150
151
152
153
154
155
156
157
158
159
160
161

    prompts = []
    sampling_params: list[SamplingParams] = []
    for request in requests:
        prompts.append(request.prompt)
        sampling_params.append(
            SamplingParams(
                n=n,
                temperature=1.0,
                top_p=1.0,
                ignore_eos=True,
                max_tokens=request.expected_output_len,
                detokenize=not disable_detokenize,
162
163
            )
        )
164
    start = time.perf_counter()
165
166
    if do_profile:
        llm.start_profile()
167
    outputs = llm.chat(prompts, sampling_params, use_tqdm=True)
168
169
    if do_profile:
        llm.stop_profile()
170
171
172
173
174
175
176
177
    end = time.perf_counter()
    return end - start, outputs


async def run_vllm_async(
    requests: list[SampleRequest],
    n: int,
    engine_args: AsyncEngineArgs,
178
    do_profile: bool,
179
180
181
182
    disable_frontend_multiprocessing: bool = False,
    disable_detokenize: bool = False,
) -> float:
    from vllm import SamplingParams
183
    from vllm.entrypoints.openai.api_server import (
184
185
        build_async_engine_client_from_engine_args,
    )
186
187

    async with build_async_engine_client_from_engine_args(
188
189
190
        engine_args,
        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
    ) as llm:
191
        model_config = llm.model_config
192
        assert all(
193
194
195
196
197
198
199
            model_config.max_model_len
            >= (request.prompt_len + request.expected_output_len)
            for request in requests
        ), (
            "Please ensure that max_model_len is greater than the sum of"
            " prompt_len and expected_output_len for all requests."
        )
200
201

        # Add the requests to the engine.
202
        prompts: list[TextPrompt | TokensPrompt] = []
203
        sampling_params: list[SamplingParams] = []
204
        lora_requests: list[LoRARequest | None] = []
205
        for request in requests:
206
207
            prompt = (
                TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"])
208
                if "prompt_token_ids" in request.prompt
209
                else TextPrompt(prompt=request.prompt)
210
            )
211
212
213
214
215

            if request.multi_modal_data:
                assert isinstance(request.multi_modal_data, dict)
                prompt["multi_modal_data"] = request.multi_modal_data

216
217
218
219
220
221
222
223
            sampling_params.append(
                SamplingParams(
                    n=n,
                    temperature=1.0,
                    top_p=1.0,
                    ignore_eos=True,
                    max_tokens=request.expected_output_len,
                    detokenize=not disable_detokenize,
224
225
                )
            )
226
            prompts.append(prompt)
227
228
229
230
            lora_requests.append(request.lora_request)

        generators = []
        start = time.perf_counter()
231
232
        if do_profile:
            await llm.start_profile()
233
234
235
236
        for i, (prompt, sp, lr) in enumerate(
            zip(prompts, sampling_params, lora_requests)
        ):
            generator = llm.generate(prompt, sp, lora_request=lr, request_id=f"test{i}")
237
238
239
240
            generators.append(generator)
        all_gens = merge_async_iterators(*generators)
        async for i, res in all_gens:
            pass
241
242
        if do_profile:
            await llm.stop_profile()
243
244
245
246
247
248
249
        end = time.perf_counter()
        return end - start


def run_hf(
    requests: list[SampleRequest],
    model: str,
250
    tokenizer: TokenizerLike,
251
252
253
254
255
    n: int,
    max_batch_size: int,
    trust_remote_code: bool,
    disable_detokenize: bool = False,
) -> float:
256
257
258
    assert isinstance(tokenizer, PreTrainedTokenizerBase), (
        "the hf backend only supports HF tokenizers"
    )
259
    llm = AutoModelForCausalLM.from_pretrained(
260
        model, dtype=torch.float16, trust_remote_code=trust_remote_code
261
    )
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
    if llm.config.model_type == "llama":
        # To enable padding in the HF backend.
        tokenizer.pad_token = tokenizer.eos_token
    llm = llm.cuda()

    pbar = tqdm(total=len(requests))
    start = time.perf_counter()
    batch: list[str] = []
    max_prompt_len = 0
    max_output_len = 0
    for i in range(len(requests)):
        prompt = requests[i].prompt
        prompt_len = requests[i].prompt_len
        output_len = requests[i].expected_output_len
        # Add the prompt to the batch.
        batch.append(prompt)
        max_prompt_len = max(max_prompt_len, prompt_len)
        max_output_len = max(max_output_len, output_len)
        if len(batch) < max_batch_size and i != len(requests) - 1:
            # Check if we can add more requests to the batch.
            next_prompt_len = requests[i + 1].prompt_len
            next_output_len = requests[i + 1].expected_output_len
284
285
286
287
            if (
                max(max_prompt_len, next_prompt_len)
                + max(max_output_len, next_output_len)
            ) <= 2048:
288
289
290
291
                # We can add more requests to the batch.
                continue

        # Generate the sequences.
292
        input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
        llm_outputs = llm.generate(
            input_ids=input_ids.cuda(),
            do_sample=True,
            num_return_sequences=n,
            temperature=1.0,
            top_p=1.0,
            use_cache=True,
            max_new_tokens=max_output_len,
        )
        if not disable_detokenize:
            # Include the decoding time.
            tokenizer.batch_decode(llm_outputs, skip_special_tokens=True)
        pbar.update(len(batch))

        # Clear the batch.
        batch = []
        max_prompt_len = 0
        max_output_len = 0
    end = time.perf_counter()
    return end - start


315
316
317
def save_to_pytorch_benchmark_format(
    args: argparse.Namespace, results: dict[str, Any]
) -> None:
318
319
320
321
322
323
324
    pt_records = convert_to_pytorch_benchmark_format(
        args=args,
        metrics={
            "requests_per_second": [results["requests_per_second"]],
            "tokens_per_second": [results["tokens_per_second"]],
        },
        extra_info={
325
326
327
            k: results[k] for k in ["elapsed_time", "num_requests", "total_num_tokens"]
        },
    )
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
    if pt_records:
        # Don't use json suffix here as we don't want CI to pick it up
        pt_file = f"{os.path.splitext(args.output_json)[0]}.pytorch.json"
        write_to_json(pt_file, pt_records)


def get_requests(args, tokenizer):
    # Common parameters for all dataset types.
    common_kwargs = {
        "dataset_path": args.dataset_path,
        "random_seed": args.seed,
    }
    sample_kwargs = {
        "tokenizer": tokenizer,
        "lora_path": args.lora_path,
        "max_loras": args.max_loras,
        "num_requests": args.num_prompts,
        "input_len": args.input_len,
        "output_len": args.output_len,
    }

349
350
351
352
    if args.dataset_name == "random" or (
        args.dataset_path is None
        and args.dataset_name not in {"prefix_repetition", "random-mm", "random-rerank"}
    ):
353
354
355
356
357
358
359
360
361
        sample_kwargs["range_ratio"] = args.random_range_ratio
        sample_kwargs["prefix_len"] = args.prefix_len
        dataset_cls = RandomDataset
    elif args.dataset_name == "sharegpt":
        dataset_cls = ShareGPTDataset
        if args.backend == "vllm-chat":
            sample_kwargs["enable_multimodal_chat"] = True
    elif args.dataset_name == "sonnet":
        assert tokenizer.chat_template or tokenizer.default_chat_template, (
362
363
            "Tokenizer/model must have chat template for sonnet dataset."
        )
364
365
366
367
368
369
370
371
        dataset_cls = SonnetDataset
        sample_kwargs["prefix_len"] = args.prefix_len
        sample_kwargs["return_prompt_formatted"] = True
    elif args.dataset_name == "burstgpt":
        dataset_cls = BurstGPTDataset
    elif args.dataset_name == "hf":
        if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
            dataset_cls = VisionArenaDataset
372
373
            common_kwargs["dataset_subset"] = None
            common_kwargs["dataset_split"] = "train"
374
375
376
            sample_kwargs["enable_multimodal_chat"] = True
        elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:
            dataset_cls = InstructCoderDataset
377
            common_kwargs["dataset_split"] = "train"
378
379
380
381
382
        elif args.dataset_path in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS:
            dataset_cls = MultiModalConversationDataset
            common_kwargs["dataset_subset"] = args.hf_subset
            common_kwargs["dataset_split"] = args.hf_split
            sample_kwargs["enable_multimodal_chat"] = True
383
384
        elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
            dataset_cls = ConversationDataset
385
386
            common_kwargs["dataset_subset"] = args.hf_subset
            common_kwargs["dataset_split"] = args.hf_split
387
388
389
            sample_kwargs["enable_multimodal_chat"] = True
        elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
            dataset_cls = AIMODataset
390
391
            common_kwargs["dataset_subset"] = None
            common_kwargs["dataset_split"] = "train"
392
393
394
395
396
397
    elif args.dataset_name == "prefix_repetition":
        dataset_cls = PrefixRepetitionRandomDataset
        sample_kwargs["prefix_len"] = args.prefix_repetition_prefix_len
        sample_kwargs["suffix_len"] = args.prefix_repetition_suffix_len
        sample_kwargs["num_prefixes"] = args.prefix_repetition_num_prefixes
        sample_kwargs["output_len"] = args.prefix_repetition_output_len
398
399
400
401
    else:
        raise ValueError(f"Unknown dataset name: {args.dataset_name}")
    # Remove None values
    sample_kwargs = {k: v for k, v in sample_kwargs.items() if v is not None}
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
    requests = dataset_cls(**common_kwargs).sample(**sample_kwargs)
    requests = filter_requests_for_dp(requests, args.data_parallel_size)
    return requests


def filter_requests_for_dp(requests, data_parallel_size):
    # Note(zhuohan): The way we get data_parallel_rank is hacky and only
    # works for external launcher mode. Should be cleaned up and deprecated
    # in the future with a better vLLM distributed process design.
    if data_parallel_size == 1:
        return requests

    global_rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    data_parallel_rank = global_rank // (world_size // data_parallel_size)
417
418
419
420
421
    return [
        r
        for i, r in enumerate(requests)
        if i % data_parallel_size == data_parallel_rank
    ]
422
423
424
425
426
427
428
429
430
431
432
433


def validate_args(args):
    """
    Validate command-line arguments.
    """

    # === Deprecation and Defaulting ===
    if args.dataset is not None:
        warnings.warn(
            "The '--dataset' argument will be deprecated in the next release. "
            "Please use '--dataset-name' and '--dataset-path' instead.",
434
435
            stacklevel=2,
        )
436
437
438
439
440
441
442
443
444
445
446
        args.dataset_path = args.dataset

    if not getattr(args, "tokenizer", None):
        args.tokenizer = args.model

    # === Backend Validation ===
    valid_backends = {"vllm", "hf", "mii", "vllm-chat"}
    if args.backend not in valid_backends:
        raise ValueError(f"Unsupported backend: {args.backend}")

    # === Dataset Configuration ===
447
448
449
450
451
    if (
        not args.dataset
        and not args.dataset_path
        and args.dataset_name not in {"prefix_repetition"}
    ):
452
453
        print("When dataset path is not set, it will default to random dataset")
        args.dataset_name = "random"
454
455
456
457
458
459
460
        if args.input_len is None:
            raise ValueError("input_len must be provided for a random dataset")

    # === Dataset Name Specific Checks ===
    # --hf-subset and --hf-split: only used
    # when dataset_name is 'hf'
    if args.dataset_name != "hf" and (
461
462
463
464
465
        getattr(args, "hf_subset", None) is not None
        or getattr(args, "hf_split", None) is not None
    ):
        warnings.warn(
            "--hf-subset and --hf-split will be ignored \
466
                since --dataset-name is not 'hf'.",
467
468
            stacklevel=2,
        )
469
470
    elif args.dataset_name == "hf":
        if args.dataset_path in (
471
            VisionArenaDataset.SUPPORTED_DATASET_PATHS.keys()
472
            | MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
473
474
475
476
            | ConversationDataset.SUPPORTED_DATASET_PATHS
        ):
            assert args.backend == "vllm-chat", (
                f"{args.dataset_path} needs to use vllm-chat as the backend."
477
            )
478
479
480
481
482
483
        elif args.dataset_path in (
            InstructCoderDataset.SUPPORTED_DATASET_PATHS
            | AIMODataset.SUPPORTED_DATASET_PATHS
        ):
            assert args.backend == "vllm", (
                f"{args.dataset_path} needs to use vllm as the backend."
484
            )
485
        else:
486
            raise ValueError(f"{args.dataset_path} is not supported by hf dataset.")
487
488

    # --random-range-ratio: only used when dataset_name is 'random'
489
490
491
    if args.dataset_name != "random" and args.random_range_ratio is not None:
        warnings.warn(
            "--random-range-ratio will be ignored since \
492
                --dataset-name is not 'random'.",
493
494
            stacklevel=2,
        )
495
496
497

    # --prefix-len: only used when dataset_name is 'random', 'sonnet', or not
    # set.
498
499
500
501
502
503
    if (
        args.dataset_name not in {"random", "sonnet", None}
        and args.prefix_len is not None
    ):
        warnings.warn(
            "--prefix-len will be ignored since --dataset-name\
504
                 is not 'random', 'sonnet', or not set.",
505
506
            stacklevel=2,
        )
507
508
509

    # === LoRA Settings ===
    if getattr(args, "enable_lora", False) and args.backend != "vllm":
510
        raise ValueError("LoRA benchmarking is only supported for vLLM backend")
511
512
513
514
515
516
517
518
519
    if getattr(args, "enable_lora", False) and args.lora_path is None:
        raise ValueError("LoRA path must be provided when enable_lora is True")

    # === Backend-specific Validations ===
    if args.backend == "hf" and args.hf_max_batch_size is None:
        raise ValueError("HF max batch size is required for HF backend")
    if args.backend != "hf" and args.hf_max_batch_size is not None:
        raise ValueError("HF max batch size is only for HF backend.")

520
521
522
523
    if (
        args.backend in {"hf", "mii"}
        and getattr(args, "quantization", None) is not None
    ):
524
525
526
527
528
529
530
        raise ValueError("Quantization is only for vLLM backend.")

    if args.backend == "mii" and args.dtype != "auto":
        raise ValueError("dtype must be auto for MII backend.")
    if args.backend == "mii" and args.n != 1:
        raise ValueError("n must be 1 for MII backend.")
    if args.backend == "mii" and args.tokenizer != args.model:
531
        raise ValueError("Tokenizer must be the same as the model for MII backend.")
532
533

    if args.data_parallel_size > 1 and (
534
535
        args.distributed_executor_backend != "external_launcher" or args.async_engine
    ):
536
537
538
539
        # --data-parallel is not supported fully.
        # Old issue: https://github.com/vllm-project/vllm/issues/16222
        # Currently we only support data parallel with external launcher
        # mode (i.e., launch with toruchrun).
540
        raise ValueError(
541
542
            "Data parallel is only supported with external launcher mode "
            "with synchronous engine in offline benchmark, "
543
544
            "please use benchmark serving instead"
        )
545
546
547


def add_cli_args(parser: argparse.ArgumentParser):
548
549
550
551
552
553
    parser.add_argument(
        "--backend",
        type=str,
        choices=["vllm", "hf", "mii", "vllm-chat"],
        default="vllm",
    )
554
555
556
    parser.add_argument(
        "--dataset-name",
        type=str,
557
        choices=["sharegpt", "random", "sonnet", "burstgpt", "hf", "prefix_repetition"],
558
        help="Name of the dataset to benchmark on.",
559
560
        default="sharegpt",
    )
561
562
563
564
565
566
567
    parser.add_argument(
        "--dataset",
        type=str,
        default=None,
        help="Path to the ShareGPT dataset, will be deprecated in\
            the next release. The dataset is expected to "
        "be a json in form of list[dict[..., conversations: "
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
        "list[dict[..., value: <prompt_or_response>]]]]",
    )
    parser.add_argument(
        "--dataset-path", type=str, default=None, help="Path to the dataset"
    )
    parser.add_argument(
        "--input-len",
        type=int,
        default=None,
        help="Input prompt length for each request",
    )
    parser.add_argument(
        "--output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the "
        "output length from the dataset.",
    )
586
    parser.add_argument(
587
588
589
590
591
592
593
594
595
596
597
598
599
        "--n", type=int, default=1, help="Number of generated sequences per prompt."
    )
    parser.add_argument(
        "--num-prompts", type=int, default=1000, help="Number of prompts to process."
    )
    parser.add_argument(
        "--hf-max-batch-size",
        type=int,
        default=None,
        help="Maximum batch size for HF backend.",
    )
    parser.add_argument(
        "--output-json",
600
601
        type=str,
        default=None,
602
603
604
605
606
607
608
609
610
611
612
613
614
615
        help="Path to save the throughput results in JSON format.",
    )
    parser.add_argument(
        "--async-engine",
        action="store_true",
        default=False,
        help="Use vLLM async engine rather than LLM class.",
    )
    parser.add_argument(
        "--disable-frontend-multiprocessing",
        action="store_true",
        default=False,
        help="Disable decoupled async engine frontend.",
    )
616
617
618
    parser.add_argument(
        "--disable-detokenize",
        action="store_true",
619
620
621
622
623
        help=(
            "Do not detokenize the response (i.e. do not include "
            "detokenization time in the measurement)"
        ),
    )
624
625
626
627
628
629
    # LoRA
    parser.add_argument(
        "--lora-path",
        type=str,
        default=None,
        help="Path to the lora adapters to use. This can be an absolute path, "
630
631
        "a relative path, or a Hugging Face model identifier.",
    )
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
    parser.add_argument(
        "--prefix-len",
        type=int,
        default=0,
        help="Number of fixed prefix tokens before the random "
        "context in a request (default: 0).",
    )
    # random dataset
    parser.add_argument(
        "--random-range-ratio",
        type=float,
        default=0.0,
        help="Range ratio for sampling input/output length, "
        "used only for RandomDataset. Must be in the range [0, 1) to define "
        "a symmetric sampling range "
        "[length * (1 - range_ratio), length * (1 + range_ratio)].",
    )

    # hf dtaset
651
652
653
654
655
656
    parser.add_argument(
        "--hf-subset", type=str, default=None, help="Subset of the HF dataset."
    )
    parser.add_argument(
        "--hf-split", type=str, default=None, help="Split of the HF dataset."
    )
657
658
659
660
    parser.add_argument(
        "--profile",
        action="store_true",
        default=False,
661
        help="Use vLLM Profiling. --profiler-config must be provided on the server.",
662
    )
663

664
665
    # prefix repetition dataset
    prefix_repetition_group = parser.add_argument_group(
666
667
        "prefix repetition dataset options"
    )
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
    prefix_repetition_group.add_argument(
        "--prefix-repetition-prefix-len",
        type=int,
        default=None,
        help="Number of prefix tokens per request, used only for prefix "
        "repetition dataset.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-suffix-len",
        type=int,
        default=None,
        help="Number of suffix tokens per request, used only for prefix "
        "repetition dataset. Total input length is prefix_len + suffix_len.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-num-prefixes",
        type=int,
        default=None,
        help="Number of prefixes to generate, used only for prefix repetition "
        "dataset. Prompts per prefix is num_requests // num_prefixes.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-output-len",
        type=int,
        default=None,
        help="Number of output tokens per request, used only for prefix "
        "repetition dataset.",
    )

697
698
699
700
701
702
703
704
705
    parser = AsyncEngineArgs.add_cli_args(parser)


def main(args: argparse.Namespace):
    validate_args(args)
    if args.seed is None:
        args.seed = 0
    random.seed(args.seed)
    # Sample the requests.
706
707
708
709
710
711
712
713
714
715
    if (
        args.backend == "hf" or args.backend == "mii"
    ) and args.tokenizer_mode == "auto":
        # mistral_common tokenizer is only supported on vllm and vllm-chat backends;
        # for hf and mii backends, we use hf tokenizer
        args.tokenizer_mode = "hf"
    tokenizer = get_tokenizer(
        args.tokenizer,
        tokenizer_mode=args.tokenizer_mode,
        trust_remote_code=args.trust_remote_code,
716
    )
717
    requests = get_requests(args, tokenizer)
718
    is_multi_modal = any(request.multi_modal_data is not None for request in requests)
719
    request_outputs: list[RequestOutput] | None = None
720
721
722
723
724
725
726
    if args.backend == "vllm":
        if args.async_engine:
            elapsed_time = uvloop.run(
                run_vllm_async(
                    requests,
                    args.n,
                    AsyncEngineArgs.from_cli_args(args),
727
728
729
                    disable_frontend_multiprocessing=args.disable_frontend_multiprocessing,
                    disable_detokenize=args.disable_detokenize,
                    do_profile=args.profile,
730
731
                )
            )
732
733
        else:
            elapsed_time, request_outputs = run_vllm(
734
735
736
                requests,
                args.n,
                EngineArgs.from_cli_args(args),
737
                disable_detokenize=args.disable_detokenize,
738
739
                do_profile=args.profile,
            )
740
741
    elif args.backend == "hf":
        assert args.tensor_parallel_size == 1
742
        if args.profile:
743
744
745
746
747
748
749
750
751
752
            raise NotImplementedError("Profiling not implemented yet for backend='hf'.")
        elapsed_time = run_hf(
            requests,
            args.model,
            tokenizer,
            args.n,
            args.hf_max_batch_size,
            args.trust_remote_code,
            args.disable_detokenize,
        )
753
754
    elif args.backend == "vllm-chat":
        elapsed_time, request_outputs = run_vllm_chat(
755
756
757
758
759
760
            requests,
            args.n,
            EngineArgs.from_cli_args(args),
            disable_detokenize=args.disable_detokenize,
            do_profile=args.profile,
        )
761
762
763
764
765
766
767
768
769
770
771
    else:
        raise ValueError(f"Unknown backend: {args.backend}")

    if request_outputs:
        # Note: with the vllm and vllm-chat backends,
        # we have request_outputs, which we use to count tokens.
        total_prompt_tokens = 0
        total_output_tokens = 0
        for ro in request_outputs:
            if not isinstance(ro, RequestOutput):
                continue
772
773
774
775
            total_prompt_tokens += (
                len(ro.prompt_token_ids) if ro.prompt_token_ids else 0
            )
            total_output_tokens += sum(len(o.token_ids) for o in ro.outputs if o)
776
777
        total_num_tokens = total_prompt_tokens + total_output_tokens
    else:
778
        total_num_tokens = sum(r.prompt_len + r.expected_output_len for r in requests)
779
780
781
782
        total_output_tokens = sum(r.expected_output_len for r in requests)
        total_prompt_tokens = total_num_tokens - total_output_tokens

    if is_multi_modal and args.backend != "vllm-chat":
783
784
785
786
787
788
        print(
            "\033[91mWARNING\033[0m: Multi-modal request with "
            f"{args.backend} backend detected. The "
            "following metrics are not accurate because image tokens are not"
            " counted. See vllm-project/vllm/issues/9778 for details."
        )
789
790
791
        # TODO(vllm-project/vllm/issues/9778): Count multi-modal token length.
        # vllm-chat backend counts the image tokens now

792
793
794
795
796
    print(
        f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "
        f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "
        f"{total_output_tokens / elapsed_time:.2f} output tokens/s"
    )
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
    print(f"Total num prompt tokens:  {total_prompt_tokens}")
    print(f"Total num output tokens:  {total_output_tokens}")

    # Output JSON results if specified
    if args.output_json:
        results = {
            "elapsed_time": elapsed_time,
            "num_requests": len(requests),
            "total_num_tokens": total_num_tokens,
            "requests_per_second": len(requests) / elapsed_time,
            "tokens_per_second": total_num_tokens / elapsed_time,
        }
        with open(args.output_json, "w") as f:
            json.dump(results, f, indent=4)
        save_to_pytorch_benchmark_format(args, results)