# 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>

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,

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,
 }
 
 

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