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 14, 05:30 UTC
Add docs for PrefixRepetitionDataset + enable usage with `vllm bench throughput` (#23012)
Signed-off-by: Seiji Eicher <seiji@anyscale.com> Co-authored-by: Roger Wang <hey@rogerw.me>
Seiji Eicher committed
Aug 16, 2025 at 03:21 UTC
de9cb617637deabab4e34db05d26c8d4d6b2ed98
2 files changed
+73
-6
benchmarks/README.md
+21
-1
@@ -40,7 +40,7 @@ become available.
40
<td><code>wget https://github.com/HPMLL/BurstGPT/releases/download/v1.1/BurstGPT_without_fails_2.csv</code></td>
41
</tr>
42
<tr>
43
- <td><strong>Sonnet</strong></td>
43
+ <td><strong>Sonnet (deprecated)</strong></td>
44
<td style="text-align: center;">✅</td>
45
<td style="text-align: center;">✅</td>
46
<td>Local file: <code>benchmarks/sonnet.txt</code></td>
@@ -51,6 +51,12 @@ become available.
51
<td style="text-align: center;">✅</td>
52
<td><code>synthetic</code></td>
53
</tr>
54
+ <tr>
55
+ <td><strong>Prefix Repetition</strong></td>
56
+ <td style="text-align: center;">✅</td>
57
+ <td style="text-align: center;">✅</td>
58
+ <td><code>synthetic</code></td>
59
+ </tr>
60
<tr>
61
<td><strong>HuggingFace-VisionArena</strong></td>
62
<td style="text-align: center;">✅</td>
@@ -592,6 +598,20 @@ python3 benchmarks/benchmark_prefix_caching.py \
598
--input-length-range 128:256
599
```
600
601
+### Prefix Repetition Dataset
602
+
603
+```bash
604
+vllm bench serve \
605
+ --backend openai \
606
+ --model meta-llama/Llama-2-7b-chat-hf \
607
+ --dataset-name prefix_repetition \
608
+ --num-prompts 100 \
609
+ --prefix-repetition-prefix-len 512 \
610
+ --prefix-repetition-suffix-len 128 \
611
+ --prefix-repetition-num-prefixes 5 \
612
+ --prefix-repetition-output-len 128
613
+```
614
+
615
</details>
616
617
## ⚡ Example - Request Prioritization Benchmark
vllm/benchmarks/throughput.py
+52
-5
@@ -18,9 +18,11 @@ from transformers import (AutoModelForCausalLM, AutoTokenizer,
18
19
from vllm.benchmarks.datasets import (AIMODataset, BurstGPTDataset,
20
ConversationDataset,
21
- InstructCoderDataset, RandomDataset,
22
- SampleRequest, ShareGPTDataset,
23
- SonnetDataset, VisionArenaDataset)
21
+ InstructCoderDataset,
22
+ PrefixRepetitionRandomDataset,
23
+ RandomDataset, SampleRequest,
24
+ ShareGPTDataset, SonnetDataset,
25
+ VisionArenaDataset)
26
from vllm.benchmarks.lib.utils import (convert_to_pytorch_benchmark_format,
27
write_to_json)
28
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
@@ -327,6 +329,12 @@ def get_requests(args, tokenizer):
329
dataset_cls = AIMODataset
330
common_kwargs['dataset_subset'] = None
331
common_kwargs['dataset_split'] = "train"
332
+ elif args.dataset_name == "prefix_repetition":
333
+ dataset_cls = PrefixRepetitionRandomDataset
334
+ sample_kwargs["prefix_len"] = args.prefix_repetition_prefix_len
335
+ sample_kwargs["suffix_len"] = args.prefix_repetition_suffix_len
336
+ sample_kwargs["num_prefixes"] = args.prefix_repetition_num_prefixes
337
+ sample_kwargs["output_len"] = args.prefix_repetition_output_len
338
else:
339
raise ValueError(f"Unknown dataset name: {args.dataset_name}")
340
# Remove None values
@@ -356,7 +364,11 @@ def validate_args(args):
364
raise ValueError(f"Unsupported backend: {args.backend}")
365
366
# === Dataset Configuration ===
359
- if not args.dataset and not args.dataset_path:
367
+ if (
368
+ not args.dataset
369
+ and not args.dataset_path
370
+ and args.dataset_name not in {"prefix_repetition"}
371
+ ):
372
print(
373
"When dataset path is not set, it will default to random dataset")
374
args.dataset_name = 'random'
@@ -432,7 +444,10 @@ def add_cli_args(parser: argparse.ArgumentParser):
444
parser.add_argument(
445
"--dataset-name",
446
type=str,
435
- choices=["sharegpt", "random", "sonnet", "burstgpt", "hf"],
447
+ choices=[
448
+ "sharegpt", "random", "sonnet", "burstgpt", "hf",
449
+ "prefix_repetition"
450
+ ],
451
help="Name of the dataset to benchmark on.",
452
default="sharegpt")
453
parser.add_argument(
@@ -521,6 +536,38 @@ def add_cli_args(parser: argparse.ArgumentParser):
536
default=None,
537
help="Split of the HF dataset.")
538
539
+ # prefix repetition dataset
540
+ prefix_repetition_group = parser.add_argument_group(
541
+ "prefix repetition dataset options")
542
+ prefix_repetition_group.add_argument(
543
+ "--prefix-repetition-prefix-len",
544
+ type=int,
545
+ default=None,
546
+ help="Number of prefix tokens per request, used only for prefix "
547
+ "repetition dataset.",
548
+ )
549
+ prefix_repetition_group.add_argument(
550
+ "--prefix-repetition-suffix-len",
551
+ type=int,
552
+ default=None,
553
+ help="Number of suffix tokens per request, used only for prefix "
554
+ "repetition dataset. Total input length is prefix_len + suffix_len.",
555
+ )
556
+ prefix_repetition_group.add_argument(
557
+ "--prefix-repetition-num-prefixes",
558
+ type=int,
559
+ default=None,
560
+ help="Number of prefixes to generate, used only for prefix repetition "
561
+ "dataset. Prompts per prefix is num_requests // num_prefixes.",
562
+ )
563
+ prefix_repetition_group.add_argument(
564
+ "--prefix-repetition-output-len",
565
+ type=int,
566
+ default=None,
567
+ help="Number of output tokens per request, used only for prefix "
568
+ "repetition dataset.",
569
+ )
570
+
571
parser = AsyncEngineArgs.add_cli_args(parser)
572
573