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, 06:00 UTC
[Spec Decode][Benchmark] Generalize spec decode offline benchmark to more methods and datasets (#18847)
Ekagra Ranjan committed
Jun 12, 2025 at 13:30 UTC
017ef648e9eaa8c70a6e32532c861aed4a2ae463
5 files changed
+403
-223
examples/offline_inference/eagle.py
+4
@@ -137,4 +137,8 @@ def main():
137
138
139
if __name__ == "__main__":
140
+ print(
141
+ "[WARNING] Use examples/offline_inference/spec_decode.py"
142
+ " instead of this script."
143
+ )
144
main()
examples/offline_inference/spec_decode.py
new
+137
@@ -0,0 +1,137 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from transformers import AutoTokenizer
4
+
5
+from vllm import LLM, SamplingParams
6
+from vllm.benchmarks.datasets import add_dataset_parser, get_samples
7
+from vllm.v1.metrics.reader import Counter, Vector
8
+
9
+try:
10
+ from vllm.utils import FlexibleArgumentParser
11
+except ImportError:
12
+ from argparse import ArgumentParser as FlexibleArgumentParser
13
+
14
+
15
+def parse_args():
16
+ parser = FlexibleArgumentParser()
17
+ add_dataset_parser(parser)
18
+ parser.add_argument(
19
+ "--dataset",
20
+ type=str,
21
+ default="./examples/data/gsm8k.jsonl",
22
+ help="downloaded from the eagle repo "
23
+ "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
24
+ )
25
+ parser.add_argument(
26
+ "--method", type=str, default="eagle", choices=["ngram", "eagle", "eagle3"]
27
+ )
28
+ parser.add_argument("--max-num-seqs", type=int, default=8)
29
+ parser.add_argument("--num-spec-tokens", type=int, default=2)
30
+ parser.add_argument("--prompt-lookup-max", type=int, default=5)
31
+ parser.add_argument("--prompt-lookup-min", type=int, default=2)
32
+ parser.add_argument("--tp", type=int, default=1)
33
+ parser.add_argument("--draft-tp", type=int, default=1)
34
+ parser.add_argument("--enforce-eager", action="store_true")
35
+ parser.add_argument("--enable-chunked-prefill", action="store_true")
36
+ parser.add_argument("--max-num-batched-tokens", type=int, default=2048)
37
+ parser.add_argument("--temp", type=float, default=0)
38
+ parser.add_argument("--top-p", type=float, default=1.0)
39
+ parser.add_argument("--top-k", type=int, default=-1)
40
+ parser.add_argument("--print-output", action="store_true")
41
+ parser.add_argument("--output-len", type=int, default=256)
42
+ return parser.parse_args()
43
+
44
+
45
+def main():
46
+ args = parse_args()
47
+ args.endpoint_type = "openai-chat"
48
+
49
+ model_dir = "meta-llama/Llama-3.1-8B-Instruct"
50
+ tokenizer = AutoTokenizer.from_pretrained(model_dir)
51
+ max_model_len = 2048
52
+
53
+ prompts = get_samples(args, tokenizer)
54
+ # add_special_tokens is False to avoid adding bos twice when using chat templates
55
+ prompt_ids = [
56
+ tokenizer.encode(prompt.prompt, add_special_tokens=False) for prompt in prompts
57
+ ]
58
+
59
+ if args.method == "eagle" or args.method == "eagle3":
60
+ if args.method == "eagle":
61
+ eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
62
+ elif args.method == "eagle3":
63
+ eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
64
+ speculative_config = {
65
+ "method": args.method,
66
+ "model": eagle_dir,
67
+ "num_speculative_tokens": args.num_spec_tokens,
68
+ "draft_tensor_parallel_size": args.draft_tp,
69
+ "max_model_len": max_model_len,
70
+ }
71
+ elif args.method == "ngram":
72
+ speculative_config = {
73
+ "method": "ngram",
74
+ "num_speculative_tokens": args.num_spec_tokens,
75
+ "prompt_lookup_max": args.prompt_lookup_max,
76
+ "prompt_lookup_min": args.prompt_lookup_min,
77
+ "max_model_len": max_model_len,
78
+ }
79
+ else:
80
+ raise ValueError(f"unknown method: {args.method}")
81
+
82
+ llm = LLM(
83
+ model=model_dir,
84
+ trust_remote_code=True,
85
+ tensor_parallel_size=args.tp,
86
+ enable_chunked_prefill=args.enable_chunked_prefill,
87
+ max_num_batched_tokens=args.max_num_batched_tokens,
88
+ enforce_eager=args.enforce_eager,
89
+ max_model_len=max_model_len,
90
+ max_num_seqs=args.max_num_seqs,
91
+ gpu_memory_utilization=0.8,
92
+ speculative_config=speculative_config,
93
+ disable_log_stats=False,
94
+ )
95
+
96
+ sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
97
+ outputs = llm.generate(prompt_token_ids=prompt_ids, sampling_params=sampling_params)
98
+
99
+ # print the generated text
100
+ if args.print_output:
101
+ for output in outputs:
102
+ print("-" * 50)
103
+ print(f"prompt: {output.prompt}")
104
+ print(f"generated text: {output.outputs[0].text}")
105
+ print("-" * 50)
106
+
107
+ try:
108
+ metrics = llm.get_metrics()
109
+ except AssertionError:
110
+ print("Metrics are not supported in the V0 engine.")
111
+ return
112
+
113
+ num_drafts = num_accepted = 0
114
+ acceptance_counts = [0] * args.num_spec_tokens
115
+ for metric in metrics:
116
+ if metric.name == "vllm:spec_decode_num_drafts":
117
+ assert isinstance(metric, Counter)
118
+ num_drafts += metric.value
119
+ elif metric.name == "vllm:spec_decode_num_accepted_tokens":
120
+ assert isinstance(metric, Counter)
121
+ num_accepted += metric.value
122
+ elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
123
+ assert isinstance(metric, Vector)
124
+ for pos in range(len(metric.values)):
125
+ acceptance_counts[pos] += metric.values[pos]
126
+
127
+ print("-" * 50)
128
+ print(f"mean acceptance length: {1 + (num_accepted / num_drafts):.2f}")
129
+ print("-" * 50)
130
+
131
+ # print acceptance at each token position
132
+ for i in range(len(acceptance_counts)):
133
+ print(f"acceptance at token {i}:{acceptance_counts[i] / num_drafts:.2f}")
134
+
135
+
136
+if __name__ == "__main__":
137
+ main()
tests/benchmarks/test_serve_cli.py
+2
@@ -31,6 +31,8 @@ def test_bench_serve(server):
31
server.host,
32
"--port",
33
str(server.port),
34
+ "--dataset-name",
35
+ "random",
36
"--random-input-len",
37
"32",
38
"--random-output-len",
vllm/benchmarks/datasets.py
+252
@@ -50,6 +50,11 @@ try:
50
except ImportError:
51
librosa = PlaceholderModule("librosa")
52
53
+try:
54
+ from vllm.utils import FlexibleArgumentParser
55
+except ImportError:
56
+ from argparse import ArgumentParser as FlexibleArgumentParser
57
+
58
logger = logging.getLogger(__name__)
59
60
# -----------------------------------------------------------------------------
@@ -458,6 +463,253 @@ class ShareGPTDataset(BenchmarkDataset):
463
return samples
464
465
466
+def add_dataset_parser(parser: FlexibleArgumentParser):
467
+ parser.add_argument("--seed", type=int, default=0)
468
+ parser.add_argument(
469
+ "--num-prompts",
470
+ type=int,
471
+ default=1000,
472
+ help="Number of prompts to process.",
473
+ )
474
+ parser.add_argument(
475
+ "--dataset-name",
476
+ type=str,
477
+ default="random",
478
+ choices=["sharegpt", "burstgpt", "sonnet", "random", "hf", "custom"],
479
+ help="Name of the dataset to benchmark on.",
480
+ )
481
+ parser.add_argument(
482
+ "--dataset-path",
483
+ type=str,
484
+ default=None,
485
+ help="Path to the sharegpt/sonnet dataset. "
486
+ "Or the huggingface dataset ID if using HF dataset.",
487
+ )
488
+
489
+ # group for dataset specific arguments
490
+ custom_group = parser.add_argument_group("custom dataset options")
491
+ custom_group.add_argument(
492
+ "--custom-output-len",
493
+ type=int,
494
+ default=256,
495
+ help=
496
+ "Number of output tokens per request, used only for custom dataset.",
497
+ )
498
+ custom_group.add_argument(
499
+ "--custom-skip-chat-template",
500
+ action="store_true",
501
+ help=
502
+ "Skip applying chat template to prompt, used only for custom dataset.",
503
+ )
504
+
505
+ sonnet_group = parser.add_argument_group("sonnet dataset options")
506
+ sonnet_group.add_argument(
507
+ "--sonnet-input-len",
508
+ type=int,
509
+ default=550,
510
+ help=
511
+ "Number of input tokens per request, used only for sonnet dataset.",
512
+ )
513
+ sonnet_group.add_argument(
514
+ "--sonnet-output-len",
515
+ type=int,
516
+ default=150,
517
+ help=
518
+ "Number of output tokens per request, used only for sonnet dataset.",
519
+ )
520
+ sonnet_group.add_argument(
521
+ "--sonnet-prefix-len",
522
+ type=int,
523
+ default=200,
524
+ help=
525
+ "Number of prefix tokens per request, used only for sonnet dataset.",
526
+ )
527
+
528
+ sharegpt_group = parser.add_argument_group("sharegpt dataset options")
529
+ sharegpt_group.add_argument(
530
+ "--sharegpt-output-len",
531
+ type=int,
532
+ default=None,
533
+ help="Output length for each request. Overrides the output length "
534
+ "from the ShareGPT dataset.",
535
+ )
536
+
537
+ random_group = parser.add_argument_group("random dataset options")
538
+ random_group.add_argument(
539
+ "--random-input-len",
540
+ type=int,
541
+ default=1024,
542
+ help=
543
+ "Number of input tokens per request, used only for random sampling.",
544
+ )
545
+ random_group.add_argument(
546
+ "--random-output-len",
547
+ type=int,
548
+ default=128,
549
+ help=
550
+ "Number of output tokens per request, used only for random sampling.",
551
+ )
552
+ random_group.add_argument(
553
+ "--random-range-ratio",
554
+ type=float,
555
+ default=0.0,
556
+ help="Range ratio for sampling input/output length, "
557
+ "used only for random sampling. Must be in the range [0, 1) to define "
558
+ "a symmetric sampling range"
559
+ "[length * (1 - range_ratio), length * (1 + range_ratio)].",
560
+ )
561
+ random_group.add_argument(
562
+ "--random-prefix-len",
563
+ type=int,
564
+ default=0,
565
+ help=("Number of fixed prefix tokens before the random context "
566
+ "in a request. "
567
+ "The total input length is the sum of `random-prefix-len` and "
568
+ "a random "
569
+ "context length sampled from [input_len * (1 - range_ratio), "
570
+ "input_len * (1 + range_ratio)]."),
571
+ )
572
+
573
+ hf_group = parser.add_argument_group("hf dataset options")
574
+ hf_group.add_argument("--hf-subset",
575
+ type=str,
576
+ default=None,
577
+ help="Subset of the HF dataset.")
578
+ hf_group.add_argument("--hf-split",
579
+ type=str,
580
+ default=None,
581
+ help="Split of the HF dataset.")
582
+ hf_group.add_argument(
583
+ "--hf-output-len",
584
+ type=int,
585
+ default=None,
586
+ help="Output length for each request. Overrides the output lengths "
587
+ "from the sampled HF dataset.",
588
+ )
589
+
590
+
591
+def get_samples(args, tokenizer) -> list[SampleRequest]:
592
+ if args.dataset_name == "custom":
593
+ dataset = CustomDataset(dataset_path=args.dataset_path)
594
+ input_requests = dataset.sample(
595
+ num_requests=args.num_prompts,
596
+ tokenizer=tokenizer,
597
+ output_len=args.custom_output_len,
598
+ skip_chat_template=args.custom_skip_chat_template,
599
+ )
600
+
601
+ elif args.dataset_name == "sonnet":
602
+ dataset = SonnetDataset(dataset_path=args.dataset_path)
603
+ # For the "sonnet" dataset, formatting depends on the backend.
604
+ if args.endpoint_type == "openai-chat":
605
+ input_requests = dataset.sample(
606
+ num_requests=args.num_prompts,
607
+ input_len=args.sonnet_input_len,
608
+ output_len=args.sonnet_output_len,
609
+ prefix_len=args.sonnet_prefix_len,
610
+ tokenizer=tokenizer,
611
+ return_prompt_formatted=False,
612
+ )
613
+ else:
614
+ assert tokenizer.chat_template or tokenizer.default_chat_template, (
615
+ "Tokenizer/model must have chat template for sonnet dataset.")
616
+ input_requests = dataset.sample(
617
+ num_requests=args.num_prompts,
618
+ input_len=args.sonnet_input_len,
619
+ output_len=args.sonnet_output_len,
620
+ prefix_len=args.sonnet_prefix_len,
621
+ tokenizer=tokenizer,
622
+ return_prompt_formatted=True,
623
+ )
624
+
625
+ elif args.dataset_name == "hf":
626
+ # all following datasets are implemented from the
627
+ # HuggingFaceDataset base class
628
+ if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
629
+ dataset_class = VisionArenaDataset
630
+ args.hf_split = "train"
631
+ args.hf_subset = None
632
+ elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:
633
+ dataset_class = InstructCoderDataset
634
+ args.hf_split = "train"
635
+ elif args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS:
636
+ dataset_class = MTBenchDataset
637
+ args.hf_split = "train"
638
+ elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
639
+ dataset_class = ConversationDataset
640
+ elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
641
+ dataset_class = AIMODataset
642
+ args.hf_split = "train"
643
+ elif args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS: # noqa: E501
644
+ dataset_class = NextEditPredictionDataset
645
+ args.hf_split = "train"
646
+ elif args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS:
647
+ dataset_class = ASRDataset
648
+ args.hf_split = "train"
649
+ else:
650
+ supported_datasets = set([
651
+ dataset_name for cls in HuggingFaceDataset.__subclasses__()
652
+ for dataset_name in cls.SUPPORTED_DATASET_PATHS
653
+ ])
654
+ raise ValueError(
655
+ f"Unsupported dataset path: {args.dataset_path}. "
656
+ "Huggingface dataset only supports dataset_path"
657
+ f" from one of following: {supported_datasets}. "
658
+ "Please consider contributing if you would "
659
+ "like to add support for additional dataset formats.")
660
+
661
+ if dataset_class.IS_MULTIMODAL and args.endpoint_type not in [
662
+ "openai-chat",
663
+ "openai-audio",
664
+ ]:
665
+ # multi-modal benchmark is only available on OpenAI Chat backend.
666
+ raise ValueError(
667
+ "Multi-modal content is only supported on 'openai-chat' and "
668
+ "'openai-audio' backend.")
669
+ input_requests = dataset_class(
670
+ dataset_path=args.dataset_path,
671
+ dataset_subset=args.hf_subset,
672
+ dataset_split=args.hf_split,
673
+ random_seed=args.seed,
674
+ ).sample(
675
+ num_requests=args.num_prompts,
676
+ tokenizer=tokenizer,
677
+ output_len=args.hf_output_len,
678
+ )
679
+
680
+ else:
681
+ # For datasets that follow a similar structure, use a mapping.
682
+ dataset_mapping = {
683
+ "sharegpt":
684
+ lambda: ShareGPTDataset(random_seed=args.seed,
685
+ dataset_path=args.dataset_path).sample(
686
+ tokenizer=tokenizer,
687
+ num_requests=args.num_prompts,
688
+ output_len=args.sharegpt_output_len,
689
+ ),
690
+ "burstgpt":
691
+ lambda: BurstGPTDataset(random_seed=args.seed,
692
+ dataset_path=args.dataset_path).
693
+ sample(tokenizer=tokenizer, num_requests=args.num_prompts),
694
+ "random":
695
+ lambda: RandomDataset(dataset_path=args.dataset_path).sample(
696
+ tokenizer=tokenizer,
697
+ num_requests=args.num_prompts,
698
+ prefix_len=args.random_prefix_len,
699
+ input_len=args.random_input_len,
700
+ output_len=args.random_output_len,
701
+ range_ratio=args.random_range_ratio,
702
+ ),
703
+ }
704
+
705
+ try:
706
+ input_requests = dataset_mapping[args.dataset_name]()
707
+ except KeyError as err:
708
+ raise ValueError(f"Unknown dataset: {args.dataset_name}") from err
709
+
710
+ return input_requests
711
+
712
+
713
# -----------------------------------------------------------------------------
714
# Custom Dataset Implementation
715
# -----------------------------------------------------------------------------
vllm/benchmarks/serve.py
+8
-223
@@ -32,12 +32,8 @@ import numpy as np
32
from tqdm.asyncio import tqdm
33
from transformers import PreTrainedTokenizerBase
34
35
-from vllm.benchmarks.datasets import (AIMODataset, ASRDataset, BurstGPTDataset,
36
- ConversationDataset, HuggingFaceDataset,
37
- InstructCoderDataset, MTBenchDataset,
38
- NextEditPredictionDataset, RandomDataset,
39
- SampleRequest, ShareGPTDataset,
40
- SonnetDataset, VisionArenaDataset)
35
+from vllm.benchmarks.datasets import (SampleRequest, add_dataset_parser,
36
+ get_samples)
37
from vllm.benchmarks.endpoint_request_func import (ASYNC_REQUEST_FUNCS,
38
OPENAI_COMPATIBLE_BACKENDS,
39
RequestFuncInput,
@@ -543,6 +539,7 @@ def save_to_pytorch_benchmark_format(args: argparse.Namespace,
539
540
541
def add_cli_args(parser: argparse.ArgumentParser):
542
+ add_dataset_parser(parser)
543
parser.add_argument(
544
"--endpoint-type",
545
type=str,
@@ -571,20 +568,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
568
default="/v1/completions",
569
help="API endpoint.",
570
)
574
- parser.add_argument(
575
- "--dataset-name",
576
- type=str,
577
- default="random",
578
- choices=["sharegpt", "burstgpt", "sonnet", "random", "hf"],
579
- help="Name of the dataset to benchmark on.",
580
- )
581
- parser.add_argument(
582
- "--dataset-path",
583
- type=str,
584
- default=None,
585
- help="Path to the sharegpt/sonnet dataset. "
586
- "Or the huggingface dataset ID if using HF dataset.",
587
- )
571
parser.add_argument(
572
"--max-concurrency",
573
type=int,
@@ -611,12 +594,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
594
"Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501
595
)
596
parser.add_argument("--use-beam-search", action="store_true")
614
- parser.add_argument(
615
- "--num-prompts",
616
- type=int,
617
- default=1000,
618
- help="Number of prompts to process.",
619
- )
597
parser.add_argument(
598
"--logprobs",
599
type=int,
@@ -648,7 +625,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
625
"bursty requests. A higher burstiness value (burstiness > 1) "
626
"results in a more uniform arrival of requests.",
627
)
651
- parser.add_argument("--seed", type=int, default=0)
628
parser.add_argument(
629
"--trust-remote-code",
630
action="store_true",
@@ -739,89 +715,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
715
"and the blog: https://hao-ai-lab.github.io/blogs/distserve",
716
)
717
742
- # group for dataset specific arguments
743
- sonnet_group = parser.add_argument_group("sonnet dataset options")
744
- sonnet_group.add_argument(
745
- "--sonnet-input-len",
746
- type=int,
747
- default=550,
748
- help=
749
- "Number of input tokens per request, used only for sonnet dataset.",
750
- )
751
- sonnet_group.add_argument(
752
- "--sonnet-output-len",
753
- type=int,
754
- default=150,
755
- help=
756
- "Number of output tokens per request, used only for sonnet dataset.",
757
- )
758
- sonnet_group.add_argument(
759
- "--sonnet-prefix-len",
760
- type=int,
761
- default=200,
762
- help=
763
- "Number of prefix tokens per request, used only for sonnet dataset.",
764
- )
765
-
766
- sharegpt_group = parser.add_argument_group("sharegpt dataset options")
767
- sharegpt_group.add_argument(
768
- "--sharegpt-output-len",
769
- type=int,
770
- default=None,
771
- help="Output length for each request. Overrides the output length "
772
- "from the ShareGPT dataset.",
773
- )
774
-
775
- random_group = parser.add_argument_group("random dataset options")
776
- random_group.add_argument(
777
- "--random-input-len",
778
- type=int,
779
- default=1024,
780
- help=
781
- "Number of input tokens per request, used only for random sampling.",
782
- )
783
- random_group.add_argument(
784
- "--random-output-len",
785
- type=int,
786
- default=128,
787
- help=
788
- "Number of output tokens per request, used only for random sampling.",
789
- )
790
- random_group.add_argument(
791
- "--random-range-ratio",
792
- type=float,
793
- default=0.0,
794
- help="Range ratio for sampling input/output length, "
795
- "used only for random sampling. Must be in the range [0, 1) to define "
796
- "a symmetric sampling range"
797
- "[length * (1 - range_ratio), length * (1 + range_ratio)].",
798
- )
799
- random_group.add_argument(
800
- "--random-prefix-len",
801
- type=int,
802
- default=0,
803
- help="Number of fixed prefix tokens before random "
804
- " context. The length range of context in a random "
805
- " request is [random-prefix-len, "
806
- " random-prefix-len + random-prefix-len * random-range-ratio).")
807
-
808
- hf_group = parser.add_argument_group("hf dataset options")
809
- hf_group.add_argument("--hf-subset",
810
- type=str,
811
- default=None,
812
- help="Subset of the HF dataset.")
813
- hf_group.add_argument("--hf-split",
814
- type=str,
815
- default=None,
816
- help="Split of the HF dataset.")
817
- hf_group.add_argument(
818
- "--hf-output-len",
819
- type=int,
820
- default=None,
821
- help="Output length for each request. Overrides the output lengths "
822
- "from the sampled HF dataset.",
823
- )
824
-
718
sampling_group = parser.add_argument_group("sampling parameters")
719
sampling_group.add_argument(
720
"--top-p",
@@ -884,7 +777,6 @@ def main(args: argparse.Namespace):
777
random.seed(args.seed)
778
np.random.seed(args.seed)
779
887
- endpoint_type = args.endpoint_type
780
label = args.label
781
model_id = args.model
782
model_name = args.served_model_name
@@ -907,115 +799,8 @@ def main(args: argparse.Namespace):
799
"Please specify '--dataset-name' and the corresponding "
800
"'--dataset-path' if required.")
801
910
- if args.dataset_name == "sonnet":
911
- dataset = SonnetDataset(dataset_path=args.dataset_path)
912
- # For the "sonnet" dataset, formatting depends on the backend.
913
- if args.backend == "openai-chat":
914
- input_requests = dataset.sample(
915
- num_requests=args.num_prompts,
916
- input_len=args.sonnet_input_len,
917
- output_len=args.sonnet_output_len,
918
- prefix_len=args.sonnet_prefix_len,
919
- tokenizer=tokenizer,
920
- return_prompt_formatted=False,
921
- )
922
- else:
923
- assert tokenizer.chat_template or tokenizer.default_chat_template, (
924
- "Tokenizer/model must have chat template for sonnet dataset.")
925
- input_requests = dataset.sample(
926
- num_requests=args.num_prompts,
927
- input_len=args.sonnet_input_len,
928
- output_len=args.sonnet_output_len,
929
- prefix_len=args.sonnet_prefix_len,
930
- tokenizer=tokenizer,
931
- return_prompt_formatted=True,
932
- )
933
-
934
- elif args.dataset_name == "hf":
935
- # all following datasets are implemented from the
936
- # HuggingFaceDataset base class
937
- if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
938
- dataset_class = VisionArenaDataset
939
- args.hf_split = "train"
940
- args.hf_subset = None
941
- elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:
942
- dataset_class = InstructCoderDataset
943
- args.hf_split = "train"
944
- elif args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS:
945
- dataset_class = MTBenchDataset
946
- args.hf_split = "train"
947
- elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
948
- dataset_class = ConversationDataset
949
- args.hf_split = "train"
950
- elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
951
- dataset_class = AIMODataset
952
- args.hf_split = "train"
953
- elif args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS: # noqa: E501
954
- dataset_class = NextEditPredictionDataset
955
- args.hf_split = "train"
956
- elif args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS:
957
- dataset_class = ASRDataset
958
- args.hf_split = "train"
959
- else:
960
- supported_datasets = set([
961
- dataset_name for cls in HuggingFaceDataset.__subclasses__()
962
- for dataset_name in cls.SUPPORTED_DATASET_PATHS
963
- ])
964
- raise ValueError(
965
- f"Unsupported dataset path: {args.dataset_path}. "
966
- "Huggingface dataset only supports dataset_path"
967
- f" from one of following: {supported_datasets}. "
968
- "Please consider contributing if you would "
969
- "like to add support for additional dataset formats.")
970
-
971
- if dataset_class.IS_MULTIMODAL and endpoint_type not in [
972
- "openai-chat",
973
- "openai-audio",
974
- ]:
975
- # multi-modal benchmark is only available on OpenAI Chat backend.
976
- raise ValueError(
977
- "Multi-modal content is only supported on 'openai-chat' and "
978
- "'openai-audio' backend.")
979
- input_requests = dataset_class(
980
- dataset_path=args.dataset_path,
981
- dataset_subset=args.hf_subset,
982
- dataset_split=args.hf_split,
983
- random_seed=args.seed,
984
- ).sample(
985
- num_requests=args.num_prompts,
986
- tokenizer=tokenizer,
987
- output_len=args.hf_output_len,
988
- )
989
-
990
- else:
991
- # For datasets that follow a similar structure, use a mapping.
992
- dataset_mapping = {
993
- "sharegpt":
994
- lambda: ShareGPTDataset(random_seed=args.seed,
995
- dataset_path=args.dataset_path).sample(
996
- tokenizer=tokenizer,
997
- num_requests=args.num_prompts,
998
- output_len=args.sharegpt_output_len,
999
- ),
1000
- "burstgpt":
1001
- lambda: BurstGPTDataset(random_seed=args.seed,
1002
- dataset_path=args.dataset_path).
1003
- sample(tokenizer=tokenizer, num_requests=args.num_prompts),
1004
- "random":
1005
- lambda: RandomDataset(dataset_path=args.dataset_path).sample(
1006
- tokenizer=tokenizer,
1007
- num_requests=args.num_prompts,
1008
- prefix_len=args.random_prefix_len,
1009
- input_len=args.random_input_len,
1010
- output_len=args.random_output_len,
1011
- range_ratio=args.random_range_ratio,
1012
- ),
1013
- }
1014
-
1015
- try:
1016
- input_requests = dataset_mapping[args.dataset_name]()
1017
- except KeyError as err:
1018
- raise ValueError(f"Unknown dataset: {args.dataset_name}") from err
802
+ # Load the dataset.
803
+ input_requests = get_samples(args, tokenizer)
804
goodput_config_dict = check_goodput_args(args)
805
806
# Collect the sampling parameters.
@@ -1043,7 +828,7 @@ def main(args: argparse.Namespace):
828
829
benchmark_result = asyncio.run(
830
benchmark(
1046
- endpoint_type=endpoint_type,
831
+ endpoint_type=args.endpoint_type,
832
api_url=api_url,
833
base_url=base_url,
834
model_id=model_id,
@@ -1073,7 +858,7 @@ def main(args: argparse.Namespace):
858
# Setup
859
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
860
result_json["date"] = current_dt
1076
- result_json["endpoint_type"] = endpoint_type
861
+ result_json["endpoint_type"] = args.endpoint_type
862
result_json["label"] = label
863
result_json["model_id"] = model_id
864
result_json["tokenizer_id"] = tokenizer_id
@@ -1118,7 +903,7 @@ def main(args: argparse.Namespace):
903
base_model_id = model_id.split("/")[-1]
904
max_concurrency_str = (f"-concurrency{args.max_concurrency}"
905
if args.max_concurrency is not None else "")
1121
- label = label or endpoint_type
906
+ label = label or args.endpoint_type
907
file_name = f"{label}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" #noqa
908
if args.result_filename:
909
file_name = args.result_filename