# 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 <khoa.pham@radixark.ai>
2026-05-28 14:37:00 -07:00
parent cd65be98df
commit 97d129f8c6
5 changed files with 262 additions and 0 deletions
@@ -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 <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.
@@ -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,
@@ -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,
 }


@@ -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
@@ -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):