chore: import upstream snapshot with attribution

python/sglang/benchmark/datasets/speed_bench.py

@@ -0,0 +1,102 @@
+"""SPEED-Bench (nvidia/SPEED-Bench) dataset for the SGLang serving benchmark.
+
+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