From 97d129f8c6e6b24ca6cfd24f3b4a154d9d339fa8 Mon Sep 17 00:00:00 2001 From: Jonathan Mamou Date: Fri, 29 May 2026 00:37:00 +0300 Subject: [PATCH] # feat(bench): add SPEED-Bench dataset support to bench_serving (#24149) Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com> Co-authored-by: Khoa Pham --- .../docs/developer_guide/bench_serving.mdx | 24 ++++ python/sglang/bench_serving.py | 14 ++ python/sglang/benchmark/datasets/__init__.py | 2 + .../sglang/benchmark/datasets/speed_bench.py | 102 +++++++++++++++ .../bench_fn/test_benchmark_datasets_api.py | 120 ++++++++++++++++++ 5 files changed, 262 insertions(+) create mode 100644 python/sglang/benchmark/datasets/speed_bench.py diff --git a/docs_new/docs/developer_guide/bench_serving.mdx b/docs_new/docs/developer_guide/bench_serving.mdx index b77097808..4513572b4 100644 --- a/docs_new/docs/developer_guide/bench_serving.mdx +++ b/docs_new/docs/developer_guide/bench_serving.mdx @@ -65,6 +65,7 @@ Select with `--dataset-name`: - `image`: generates images and wraps them in chat messages; supports custom resolutions, multiple formats, and different content types - `generated-shared-prefix`: synthetic dataset with shared long system prompts and short questions - `mmmu`: samples from MMMU (Math split) and includes images +- `speed-bench`: [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) (**SPEculative Evaluation Dataset**) — a unified benchmark for evaluating [Speculative Decoding (SD)](https://arxiv.org/abs/2604.09557) algorithms. Uses the Throughput split, which provides fixed-length input sequences (1K–32K tokens) grouped into three output-entropy categories (`low_entropy`, `mixed`, `high_entropy`). Requires a pre-downloaded JSONL file passed via `--dataset-path`. Common dataset flags: @@ -90,6 +91,12 @@ Image dataset flags (for `image`): - `--image-format`: Image format (jpeg or png) - `--image-content`: Image content type (random or blank) +SPEED-Bench flags (for `speed-bench`): + +- `--dataset-path PATH`: path to the pre-downloaded SPEED-Bench Throughput JSONL (e.g., `throughput_1k.jsonl`). Use the [SPEED-Bench measurement framework](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/specdec_bench) to generate it. +- `--speed-bench-category`: filter to one entropy category: `low_entropy`, `mixed`, or `high_entropy` (default: all) +- `--speed-bench-output-len`: fixed number of output tokens per request (default: 512) + ### Examples 1. To benchmark image dataset with 3 images per request, 500 prompts, 512 input length, and 512 output length, you can run: @@ -125,6 +132,23 @@ python3 -m sglang.bench_serving \ --random-range-ratio 0.5 ``` +3. To benchmark speculative decoding throughput using SPEED-Bench (mixed-entropy category, 1K ISL), you can run: + +```bash Command +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \ + --speculative-algorithm EAGLE --speculative-draft-model-path +``` + +```bash Command +python3 -m sglang.bench_serving \ + --backend sglang \ + --dataset-name speed-bench \ + --dataset-path /path/to/throughput_1k.jsonl \ + --speed-bench-category mixed \ + --speed-bench-output-len 512 \ + --num-prompts 512 +``` + ### Choosing model and tokenizer - `--model` is required unless the backend exposes `GET /v1/models`, in which case the first model ID is auto-selected. diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index 767ad3e71..53e4ce18f 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -1982,12 +1982,26 @@ if __name__ == "__main__": "image", "mooncake", "longbench_v2", + "speed-bench", ], help="Name of the dataset to benchmark on.", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) + parser.add_argument( + "--speed-bench-category", + type=str, + default=None, + choices=["low_entropy", "mixed", "high_entropy"], + help="Category filter for the speed-bench dataset.", + ) + parser.add_argument( + "--speed-bench-output-len", + type=int, + default=512, + help="Fixed output length for speed-bench requests (default: 512).", + ) parser.add_argument( "--model", type=str, diff --git a/python/sglang/benchmark/datasets/__init__.py b/python/sglang/benchmark/datasets/__init__.py index 2320bb7df..36a437bfa 100644 --- a/python/sglang/benchmark/datasets/__init__.py +++ b/python/sglang/benchmark/datasets/__init__.py @@ -13,6 +13,7 @@ from sglang.benchmark.datasets.mooncake import MooncakeDataset from sglang.benchmark.datasets.openai_dataset import OpenAIDataset from sglang.benchmark.datasets.random import RandomDataset from sglang.benchmark.datasets.sharegpt import ShareGPTDataset +from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset DATASET_MAPPING: Dict[str, Type[BaseDataset]] = { "autobench": AutoBenchmarkDataset, @@ -28,6 +29,7 @@ DATASET_MAPPING: Dict[str, Type[BaseDataset]] = { "image": ImageDataset, "mooncake": MooncakeDataset, "longbench_v2": LongBenchV2Dataset, + "speed-bench": SpeedBenchDataset, } diff --git a/python/sglang/benchmark/datasets/speed_bench.py b/python/sglang/benchmark/datasets/speed_bench.py new file mode 100644 index 000000000..df4c9e343 --- /dev/null +++ b/python/sglang/benchmark/datasets/speed_bench.py @@ -0,0 +1,102 @@ +"""SPEED-Bench (nvidia/SPEED-Bench) dataset for SGLang bench_serving. + +Reads the pre-downloaded throughput_1k JSONL produced by prepare_speed_bench.sh +(or equivalent), optionally filtering by category (low_entropy / mixed / +high_entropy) and fixing the output length. + +CLI args consumed: + --dataset-path Path to the local JSONL file. + --speed-bench-category Category filter: low_entropy | mixed | high_entropy + (default: all categories). + --speed-bench-output-len Fixed number of output tokens per request (default: 512). + --num-prompts Number of requests to sample (capped by available rows). +""" + +import json +import random +from argparse import Namespace +from dataclasses import dataclass +from typing import List, Optional + +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import BaseDataset, DatasetRow + + +@dataclass +class SpeedBenchDataset(BaseDataset): + dataset_path: str + category: Optional[str] + output_len: int + num_requests: int + + @classmethod + def from_args(cls, args: Namespace) -> "SpeedBenchDataset": + if not args.dataset_path: + raise ValueError( + "--dataset-path must point to the SPEED-Bench JSONL file " + "(run prepare_speed_bench.sh to generate it)." + ) + return cls( + dataset_path=args.dataset_path, + category=getattr(args, "speed_bench_category", None) or None, + output_len=getattr(args, "speed_bench_output_len", 512), + num_requests=args.num_prompts, + ) + + def load( + self, tokenizer: PreTrainedTokenizerBase, model_id=None + ) -> List[DatasetRow]: + unique_prompts = [] + with open(self.dataset_path, encoding="utf-8") as f: + for line in f: + row = json.loads(line) + if self.category and row.get("category") != self.category: + continue + # turns is a list of strings; use the first user turn as the prompt + turns = row.get("turns", []) + if not turns: + continue + unique_prompts.append(turns[0]) + + if not unique_prompts: + raise ValueError( + f"No rows found in {self.dataset_path}" + + (f" for category={self.category}" if self.category else "") + ) + + # Tokenize unique prompts once to avoid redundant work + unique_dataset_rows: List[DatasetRow] = [] + for prompt_text in unique_prompts: + # Apply chat template to match vllm bench behaviour + try: + prompt_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt_text}], + add_generation_prompt=True, + tokenize=True, + ) + prompt = tokenizer.decode(prompt_ids) + except Exception: + prompt_ids = tokenizer.encode(prompt_text) + prompt = prompt_text + + unique_dataset_rows.append( + DatasetRow( + prompt=prompt, + prompt_len=len(prompt_ids), + output_len=self.output_len, + ) + ) + + # Sample (with replacement if needed); shuffle oversampled rows for + # a realistic request distribution + if self.num_requests <= len(unique_dataset_rows): + dataset_rows = random.sample(unique_dataset_rows, self.num_requests) + else: + dataset_rows = unique_dataset_rows * ( + self.num_requests // len(unique_dataset_rows) + 1 + ) + dataset_rows = dataset_rows[: self.num_requests] + random.shuffle(dataset_rows) + + return dataset_rows diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index 852a290f6..777bb6741 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/test/registered/bench_fn/test_benchmark_datasets_api.py @@ -135,6 +135,8 @@ def make_args(**overrides): "gsp_ordered": False, "seed": 1, "mooncake_workload": "conversation", + "speed_bench_category": None, + "speed_bench_output_len": 512, } args.update(overrides) return SimpleNamespace(**args) @@ -216,6 +218,39 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): f.write(json.dumps(row) + "\n") return str(path) + def _write_speed_bench_jsonl(self): + rows = [ + { + "question_id": "sb_001", + "category": "low_entropy", + "turns": ["Complete this Python function: def add(a, b):"], + }, + { + "question_id": "sb_002", + "category": "mixed", + "turns": [ + "Explain the concept of attention mechanisms in transformers." + ], + }, + { + "question_id": "sb_003", + "category": "high_entropy", + "turns": ["Write a short story about a robot discovering music."], + }, + { + "question_id": "sb_004", + "category": "low_entropy", + "turns": [ + "Sort the following list in ascending order: [5, 2, 8, 1, 9]" + ], + }, + ] + path = self.tmpdir_path / "speed_bench.jsonl" + with open(path, "w") as f: + for row in rows: + f.write(json.dumps(row) + "\n") + return str(path) + def _write_mooncake_jsonl(self): rows = [ {"timestamp": 1000, "hash_ids": [1, 2], "output_length": 5}, @@ -356,6 +391,78 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): self.assertEqual(len(rows), 2) self.assertTrue(all(isinstance(row, DatasetRow) for row in rows)) + def test_speed_bench_sampler(self): + dataset_path = self._write_speed_bench_jsonl() + args = make_args( + dataset_name="speed-bench", + dataset_path=dataset_path, + num_prompts=3, + ) + from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset + + dataset = SpeedBenchDataset.from_args(args) + rows = dataset.load(self.tokenizer) + self.assertEqual(len(rows), 3) + self.assertTrue(all(isinstance(row, DatasetRow) for row in rows)) + self.assertTrue(all(row.output_len == 512 for row in rows)) + self.assertTrue(all(row.prompt_len > 0 for row in rows)) + + def test_speed_bench_category_filter(self): + dataset_path = self._write_speed_bench_jsonl() + args = make_args( + dataset_name="speed-bench", + dataset_path=dataset_path, + num_prompts=2, + speed_bench_category="low_entropy", + ) + from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset + + dataset = SpeedBenchDataset.from_args(args) + rows = dataset.load(self.tokenizer) + # Only 2 low_entropy rows in the fixture, num_prompts=2 + self.assertEqual(len(rows), 2) + self.assertTrue(all(isinstance(row, DatasetRow) for row in rows)) + + def test_speed_bench_output_len_override(self): + dataset_path = self._write_speed_bench_jsonl() + args = make_args( + dataset_name="speed-bench", + dataset_path=dataset_path, + num_prompts=2, + speed_bench_output_len=128, + ) + from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset + + dataset = SpeedBenchDataset.from_args(args) + rows = dataset.load(self.tokenizer) + self.assertEqual(len(rows), 2) + self.assertTrue(all(row.output_len == 128 for row in rows)) + + def test_speed_bench_empty_category_raises(self): + dataset_path = self._write_speed_bench_jsonl() + args = make_args( + dataset_name="speed-bench", + dataset_path=dataset_path, + num_prompts=1, + speed_bench_category="nonexistent_category", + ) + from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset + + dataset = SpeedBenchDataset.from_args(args) + with self.assertRaises(ValueError): + dataset.load(self.tokenizer) + + def test_speed_bench_no_path_raises(self): + args = make_args( + dataset_name="speed-bench", + dataset_path="", + num_prompts=1, + ) + from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset + + with self.assertRaises(ValueError): + SpeedBenchDataset.from_args(args) + def test_dataset_mapping_and_dispatch(self): expected = { "sharegpt", @@ -367,6 +474,7 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): "mmmu", "image", "mooncake", + "speed-bench", } self.assertTrue(expected.issubset(set(DATASET_MAPPING.keys()))) @@ -429,6 +537,18 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): self.assertEqual(len(gsp_rows), 4) self.assertTrue(all(isinstance(row, DatasetRow) for row in gsp_rows)) + speed_bench_path = self._write_speed_bench_jsonl() + speed_bench_args = make_args( + dataset_name="speed-bench", + dataset_path=speed_bench_path, + num_prompts=2, + ) + speed_bench_rows = get_dataset( + speed_bench_args, self.tokenizer, model_id="dummy-model" + ) + self.assertEqual(len(speed_bench_rows), 2) + self.assertTrue(all(isinstance(row, DatasetRow) for row in speed_bench_rows)) + def test_get_dataset_unknown_dataset(self): args = make_args(dataset_name="not-a-dataset") with self.assertRaises(ValueError):