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 5, 03:30 UTC
[SpecDecode][Benchmark] Add SPEED-bench support to benchmarking CLI (#36029)
Signed-off-by: talora <talora@nvidia.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Talor Abramovich committed
Apr 15, 2026 at 19:00 UTC
3cc328a4be4976f75ce016f60bc55beee4701d1b
2 files changed
+152
docs/benchmarking/cli.md
+64
@@ -37,6 +37,7 @@ th {
37
| HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
38
| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` |
39
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
40
+| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
41
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
42
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
43
@@ -239,6 +240,69 @@ vllm bench serve \
240
--spec-bench-category "summarization"
241
```
242
243
+#### SPEED-Bench Benchmark with Speculative Decoding
244
+
245
+[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
246
+
247
+!!! note
248
+ This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
249
+
250
+First, download the dataset to a folder, using this one liner:
251
+
252
+```bash
253
+curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
254
+```
255
+
256
+The command supports also the following arguments:
257
+
258
+- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
259
+- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
260
+
261
+Start a server with speculative decoding:
262
+
263
+```bash
264
+vllm serve meta-llama/Llama-3.3-70B-Instruct \
265
+ --speculative-config $'{"method": "eagle3",
266
+ "num_speculative_tokens": 3,
267
+ "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
268
+```
269
+
270
+Run all categories in the Qualitative split:
271
+
272
+```bash
273
+vllm bench serve \
274
+ --model meta-llama/Llama-3.3-70B-Instruct \
275
+ --dataset-name speed_bench \
276
+ --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
277
+ --num-prompts -1
278
+```
279
+
280
+Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
281
+
282
+Run only a specific category like "multilingual":
283
+
284
+```bash
285
+vllm bench serve \
286
+ --model meta-llama/Llama-3.3-70B-Instruct \
287
+ --dataset-name speed_bench \
288
+ --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
289
+ --num-prompts -1
290
+ --speed-bench-category "multilingual"
291
+```
292
+
293
+Run all categories in the Throughput split (2k ISL):
294
+
295
+```bash
296
+vllm bench serve \
297
+ --model meta-llama/Llama-3.3-70B-Instruct \
298
+ --dataset-name speed_bench \
299
+ --speed-bench-dataset-subset throughput_2k
300
+ --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
301
+ --num-prompts -1
302
+```
303
+
304
+Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
305
+
306
#### Other HuggingFaceDataset Examples
307
308
```bash
vllm/benchmarks/datasets/datasets.py
+88
@@ -25,6 +25,7 @@ from contextlib import suppress
25
from dataclasses import dataclass, replace
26
from functools import cache
27
from io import BytesIO
28
+from pathlib import Path
29
from tempfile import NamedTemporaryFile
30
from typing import Any, cast
31
@@ -1422,6 +1423,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1423
"custom_mm",
1424
"prefix_repetition",
1425
"spec_bench",
1426
+ "speed_bench",
1427
],
1428
help="Name of the dataset to benchmark on.",
1429
)
@@ -1606,6 +1608,34 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1608
"repetition dataset.",
1609
)
1610
1611
+ speed_bench_group = parser.add_argument_group("speed bench dataset options")
1612
+ speed_bench_group.add_argument(
1613
+ "--speed-bench-dataset-subset",
1614
+ type=str,
1615
+ default="qualitative",
1616
+ choices={
1617
+ "qualitative",
1618
+ "throughput_1k",
1619
+ "throughput_2k",
1620
+ "throughput_8k",
1621
+ "throughput_16k",
1622
+ "throughput_32k",
1623
+ },
1624
+ help="Subset of the SPEED-Bench dataset.",
1625
+ )
1626
+ speed_bench_group.add_argument(
1627
+ "--speed-bench-output-len",
1628
+ type=int,
1629
+ default=4096,
1630
+ help="Num of output tokens per request, used only for speed bench dataset.",
1631
+ )
1632
+ speed_bench_group.add_argument(
1633
+ "--speed-bench-category",
1634
+ type=str,
1635
+ default=None,
1636
+ help="Category for speed bench dataset. If None, use all categories.",
1637
+ )
1638
+
1639
1640
def add_random_dataset_base_args(
1641
parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
@@ -2074,6 +2104,19 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2104
request_id_prefix=args.request_id_prefix,
2105
no_oversample=args.no_oversample,
2106
),
2107
+ "speed_bench": lambda: SpeedBench(
2108
+ dataset_path=args.dataset_path,
2109
+ dataset_subset=args.speed_bench_dataset_subset,
2110
+ category=args.speed_bench_category,
2111
+ disable_shuffle=args.disable_shuffle,
2112
+ ).sample(
2113
+ num_requests=args.num_prompts,
2114
+ tokenizer=tokenizer,
2115
+ output_len=args.speed_bench_output_len,
2116
+ enable_multimodal_chat=args.enable_multimodal_chat,
2117
+ request_id_prefix=args.request_id_prefix,
2118
+ no_oversample=args.no_oversample,
2119
+ ),
2120
}
2121
2122
try:
@@ -3551,3 +3594,48 @@ class MMStarDataset(HuggingFaceDataset):
3594
sampled_requests, num_requests, request_id_prefix, no_oversample
3595
)
3596
return sampled_requests
3597
+
3598
+
3599
+# -----------------------------------------------------------------------------
3600
+# Speed Bench Dataset Implementation
3601
+# -----------------------------------------------------------------------------
3602
+
3603
+
3604
+class SpeedBench(CustomDataset):
3605
+ """
3606
+ Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
3607
+ Download the dataset using:
3608
+ curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -
3609
+ """ # noqa: E501
3610
+
3611
+ def __init__(self, **kwargs) -> None:
3612
+ self.dataset_subset = kwargs.pop("dataset_subset", "qualitative")
3613
+ self.category = kwargs.pop("category", None)
3614
+ super().__init__(**kwargs)
3615
+ self.load_data()
3616
+
3617
+ def load_data(self) -> None:
3618
+ if self.dataset_path is None:
3619
+ raise ValueError("dataset_path must be provided for loading data.")
3620
+
3621
+ self.data = []
3622
+
3623
+ # Load the JSONL file
3624
+ jsonl_data = pd.read_json(
3625
+ path_or_buf=Path(self.dataset_path) / f"{self.dataset_subset}.jsonl",
3626
+ lines=True,
3627
+ )
3628
+
3629
+ # check if the JSONL file has a 'turns' column
3630
+ if "messages" not in jsonl_data.columns:
3631
+ raise ValueError("JSONL file must contain a 'messages' column.")
3632
+
3633
+ for _, row in jsonl_data.iterrows():
3634
+ # sample only from a specific category if specified
3635
+ if (not self.category) or (self.category == row["category"]):
3636
+ prompt = row["messages"][0]["content"]
3637
+ self.data.append({"prompt": prompt})
3638
+
3639
+ random.seed(self.random_seed)
3640
+ if not getattr(self, "disable_shuffle", False):
3641
+ random.shuffle(self.data)