chore: import upstream snapshot with attribution

python/sglang/benchmark/datasets/__init__.py

@@ -0,0 +1,55 @@
+from typing import Dict, Type
+
+from sglang.benchmark.datasets.agentic_trace import AgenticTraceDataset
+from sglang.benchmark.datasets.autobench import AutoBenchmarkDataset
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+from sglang.benchmark.datasets.custom import CustomDataset
+from sglang.benchmark.datasets.generated_shared_prefix import (
+    GeneratedSharedPrefixDataset,
+)
+from sglang.benchmark.datasets.image import ImageDataset
+from sglang.benchmark.datasets.longbench_v2 import LongBenchV2Dataset
+from sglang.benchmark.datasets.mmmu import MMMUDataset
+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]] = {
+    "agentic-trace": AgenticTraceDataset,
+    "autobench": AutoBenchmarkDataset,
+    "sharegpt": ShareGPTDataset,
+    "custom": CustomDataset,
+    "openai": OpenAIDataset,
+    # TODO: "random" vs "random-ids" should be a flag (e.g. --random-source=sharegpt|integers),
+    # not two separate dataset names sharing the same class.
+    "random": RandomDataset,
+    "random-ids": RandomDataset,
+    "generated-shared-prefix": GeneratedSharedPrefixDataset,
+    "mmmu": MMMUDataset,
+    "image": ImageDataset,
+    "mooncake": MooncakeDataset,
+    "longbench_v2": LongBenchV2Dataset,
+    "speed-bench": SpeedBenchDataset,
+}
+
+
+def get_dataset(args, tokenizer, model_id=None):
+    dataset_name = args.dataset_name
+    if dataset_name.startswith("random") and dataset_name not in DATASET_MAPPING:
+        dataset_name = "random-ids"
+
+    if dataset_name not in DATASET_MAPPING:
+        raise ValueError(f"Unknown dataset: {args.dataset_name}")
+
+    dataset_cls = DATASET_MAPPING[dataset_name]
+    dataset = dataset_cls.from_args(args)
+    return dataset.load(tokenizer=tokenizer, model_id=model_id)
+
+
+__all__ = [
+    "DATASET_MAPPING",
+    "DatasetRow",
+    "get_dataset",
+]

python/sglang/benchmark/datasets/agentic_trace.py

@@ -0,0 +1,114 @@
+import json
+import os
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+
+# Per-turn output length when --sharegpt-output-len is not given; matches the
+# ~220-token average assistant reply of OpenHands-style agentic traces.
+DEFAULT_AGENTIC_OUTPUT_LEN = 220
+
+
+@dataclass
+class AgenticTraceDataset(BaseDataset):
+    """Multi-turn agentic trace loader (e.g. OpenHands / SWE-smith traces).
+
+    Expects a trace JSON of the shape::
+
+        {
+          "metadata": {...},
+          "conversations": [
+            [   # one conversation == a list of turns
+              {"messages": [{"role": "system", ...}, {"role": "user", ...}],
+               "prompt_tokens": 73821},
+              {"messages": [{"role": "user", ...}], "prompt_tokens": 74894},
+              ...
+            ],
+            ...
+          ]
+        }
+
+    Each turn's ``messages`` holds only the new non-assistant messages for that
+    turn. One conversation becomes one :class:`DatasetRow` whose ``prompt`` is
+    the list of per-turn message deltas; ``bench_serving`` detects this shape as
+    multi-turn and replays each conversation round by round, feeding the
+    server's real assistant reply back into the next round's history.
+
+    Use with a chat backend (``--backend sglang-oai-chat``).
+    """
+
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+    offset: int
+    max_turns: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "AgenticTraceDataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+            offset=args.dataset_offset,
+            max_turns=args.agentic_max_turns,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        if not os.path.isfile(self.dataset_path):
+            raise FileNotFoundError(f"Dataset not found at {self.dataset_path}")
+
+        with open(self.dataset_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        conversations = data.get("conversations", [])
+        if not conversations:
+            raise ValueError(f"No 'conversations' found in {self.dataset_path}.")
+
+        offset = self.offset % len(conversations)
+        if offset:
+            conversations = conversations[offset:] + conversations[:offset]
+
+        output_len = self.fixed_output_len or DEFAULT_AGENTIC_OUTPUT_LEN
+
+        filtered_dataset: List[DatasetRow] = []
+        for conversation in conversations:
+            if self.num_requests > 0 and len(filtered_dataset) >= self.num_requests:
+                break
+
+            prompt = [turn["messages"] for turn in conversation if turn.get("messages")]
+            if self.max_turns:
+                prompt = prompt[: self.max_turns]
+            if not prompt:
+                continue
+
+            # Informational only: multi-turn replay ignores per-row prompt_len.
+            prompt_len = int(conversation[0].get("prompt_tokens", 0))
+
+            filtered_dataset.append(
+                DatasetRow(
+                    prompt=prompt,
+                    prompt_len=prompt_len,
+                    output_len=output_len,
+                )
+            )
+
+        if not filtered_dataset:
+            raise ValueError(
+                f"No usable conversations loaded from {self.dataset_path}."
+            )
+
+        num_turns = [len(row.prompt) for row in filtered_dataset]
+        print(
+            f"#Conversations: {len(filtered_dataset)} "
+            f"(offset={offset}, turns/conv min={min(num_turns)} "
+            f"max={max(num_turns)} avg={np.mean(num_turns):.1f})"
+        )
+        print(f"#Output tokens per turn: {output_len}")
+        return filtered_dataset

python/sglang/benchmark/datasets/autobench.py

@@ -0,0 +1,299 @@
+import json
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+
+AUTOBENCH_RESERVED_FIELDS = {
+    "prompt",
+    "messages",
+    "prompt_origin",
+    "output_len",
+    "max_tokens",
+    "max_completion_tokens",
+    "completion_tokens",
+    "prompt_len",
+    "text_prompt_len",
+    "vision_prompt_len",
+    "image_data",
+    "timestamp",
+    "routing_key",
+    "metadata",
+    "extra_request_body",
+    "param_send",
+}
+
+
+def _load_json_if_needed(value: Any) -> Any:
+    if not isinstance(value, str):
+        return value
+    value = value.strip()
+    if not value:
+        return value
+    if value[0] not in "[{":
+        return value
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError:
+        return value
+
+
+def _normalize_messages(messages: Any) -> Optional[List[Dict[str, Any]]]:
+    messages = _load_json_if_needed(messages)
+    if not isinstance(messages, list) or not messages:
+        return None
+    if not all(isinstance(message, dict) for message in messages):
+        return None
+
+    normalized = []
+    for message in messages:
+        if "role" not in message:
+            return None
+        content = message.get("content")
+        if content is None:
+            return None
+        normalized.append({"role": message["role"], "content": content})
+    return normalized
+
+
+def _normalize_legacy_system_content(
+    system_prompt: Any, content_list: Any
+) -> Optional[List[Dict[str, Any]]]:
+    if not isinstance(content_list, list) or not content_list:
+        return None
+
+    messages: List[Dict[str, Any]] = []
+    if system_prompt:
+        messages.append({"role": "system", "content": str(system_prompt)})
+
+    turns = [str(item) for item in content_list]
+    # In the old auto_benchmark helpers, an even number of items usually means the
+    # last assistant reply is present and should be removed before benchmarking.
+    if len(turns) % 2 == 0:
+        turns = turns[:-1]
+    if not turns:
+        return None
+
+    for index, turn in enumerate(turns):
+        role = "user" if index % 2 == 0 else "assistant"
+        messages.append({"role": role, "content": turn})
+    return messages
+
+
+def _normalize_prompt(row: Dict[str, Any]) -> Tuple[Any, str]:
+    prompt = row.get("prompt")
+    messages = row.get("messages")
+    prompt_origin = row.get("prompt_origin")
+
+    if messages is not None:
+        normalized = _normalize_messages(messages)
+        if normalized is not None:
+            return normalized, "messages"
+
+    if prompt is not None:
+        prompt = _load_json_if_needed(prompt)
+        if isinstance(prompt, list) and prompt and isinstance(prompt[0], dict):
+            normalized = _normalize_messages(prompt)
+            if normalized is not None:
+                return normalized, "messages"
+        if (
+            isinstance(prompt, list)
+            and prompt
+            and all(isinstance(item, str) for item in prompt)
+        ):
+            return prompt, "multi_turn"
+        if (
+            isinstance(prompt, list)
+            and prompt
+            and all(
+                isinstance(item, list)
+                and item
+                and all(
+                    isinstance(m, dict) and "role" in m and "content" in m for m in item
+                )
+                for item in prompt
+            )
+        ):
+            # Multi-turn with N messages per round (e.g. tool observations).
+            return prompt, "multi_turn"
+        if (
+            isinstance(prompt, list)
+            and prompt
+            and all(isinstance(item, int) for item in prompt)
+        ):
+            return prompt, "token_ids"
+        if isinstance(prompt, str) and prompt:
+            return prompt, "prompt"
+
+    if prompt_origin is not None:
+        normalized = _normalize_messages(prompt_origin)
+        if normalized is not None:
+            return normalized, "messages"
+
+    if "system" in row and "content" in row:
+        normalized = _normalize_legacy_system_content(
+            row.get("system"), row.get("content")
+        )
+        if normalized is not None:
+            return normalized, "messages"
+
+    raise ValueError("Unsupported auto benchmark row: missing prompt/messages")
+
+
+def _estimate_prompt_lens(
+    prompt: Any,
+    prompt_kind: str,
+    tokenizer: PreTrainedTokenizerBase,
+    row: Dict[str, Any],
+) -> Tuple[int, int, int]:
+    if row.get("prompt_len") is not None:
+        prompt_len = int(row["prompt_len"])
+        text_prompt_len = int(row.get("text_prompt_len", prompt_len))
+        vision_prompt_len = int(row.get("vision_prompt_len", 0))
+        return prompt_len, text_prompt_len, vision_prompt_len
+
+    if prompt_kind == "messages":
+        text_prompt_len = len(
+            tokenizer.apply_chat_template(
+                prompt, tokenize=True, add_generation_prompt=True
+            )
+        )
+        vision_prompt_len = 0
+        return text_prompt_len, text_prompt_len, vision_prompt_len
+
+    if prompt_kind == "prompt":
+        prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False))
+        return prompt_len, prompt_len, 0
+
+    if prompt_kind == "token_ids":
+        prompt_len = len(prompt)
+        return prompt_len, prompt_len, 0
+
+    # Multi-turn prompt lists are handled specially by the serving benchmark and do not
+    # contribute reliable static prompt lengths.
+    return 0, 0, 0
+
+
+def _collect_extra_request_body(row: Dict[str, Any]) -> Dict[str, Any]:
+    extra: Dict[str, Any] = {}
+
+    param_send = row.get("param_send")
+    if param_send is not None:
+        parsed = _load_json_if_needed(param_send)
+        if isinstance(parsed, dict):
+            extra.update(parsed)
+
+    for key, value in row.items():
+        if key not in AUTOBENCH_RESERVED_FIELDS:
+            extra[key] = value
+
+    explicit_extra = row.get("extra_request_body")
+    explicit_extra = _load_json_if_needed(explicit_extra)
+    if isinstance(explicit_extra, dict):
+        extra.update(explicit_extra)
+
+    return extra
+
+
+def serialize_dataset_row_to_autobench(
+    row: DatasetRow, metadata: Optional[Dict[str, Any]] = None
+) -> Dict[str, Any]:
+    record: Dict[str, Any] = {
+        "prompt": row.prompt,
+        "output_len": row.output_len,
+    }
+    if row.prompt_len:
+        record["prompt_len"] = row.prompt_len
+    if row.text_prompt_len not in (None, row.prompt_len):
+        record["text_prompt_len"] = row.text_prompt_len
+    if row.vision_prompt_len:
+        record["vision_prompt_len"] = row.vision_prompt_len
+    if row.image_data:
+        record["image_data"] = row.image_data
+    if row.timestamp is not None:
+        record["timestamp"] = row.timestamp
+    if row.routing_key is not None:
+        record["routing_key"] = row.routing_key
+    if row.extra_request_body:
+        record["extra_request_body"] = row.extra_request_body
+    if metadata:
+        record["metadata"] = metadata
+    return record
+
+
+@dataclass
+class AutoBenchmarkDataset(BaseDataset):
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "AutoBenchmarkDataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_autobench_requests(
+            dataset_path=self.dataset_path,
+            num_requests=self.num_requests,
+            tokenizer=tokenizer,
+            fixed_output_len=self.fixed_output_len,
+        )
+
+
+def sample_autobench_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: PreTrainedTokenizerBase,
+    fixed_output_len: Optional[int] = None,
+) -> List[DatasetRow]:
+    dataset: List[DatasetRow] = []
+
+    with open(dataset_path, "r", encoding="utf-8") as f:
+        for line in f:
+            if num_requests > 0 and len(dataset) >= num_requests:
+                break
+
+            line = line.strip()
+            if not line:
+                continue
+
+            row = json.loads(line)
+            prompt, prompt_kind = _normalize_prompt(row)
+            prompt_len, text_prompt_len, vision_prompt_len = _estimate_prompt_lens(
+                prompt, prompt_kind, tokenizer, row
+            )
+
+            output_len = fixed_output_len or row.get("output_len")
+            output_len = output_len or row.get("max_tokens")
+            output_len = output_len or row.get("max_completion_tokens")
+            output_len = output_len or row.get("completion_tokens")
+            output_len = int(output_len or 256)
+
+            dataset.append(
+                DatasetRow(
+                    prompt=prompt,
+                    prompt_len=prompt_len,
+                    output_len=output_len,
+                    text_prompt_len=text_prompt_len,
+                    vision_prompt_len=vision_prompt_len,
+                    image_data=row.get("image_data"),
+                    timestamp=row.get("timestamp"),
+                    routing_key=row.get("routing_key"),
+                    extra_request_body=_collect_extra_request_body(row),
+                )
+            )
+
+    print(f"Loaded {len(dataset)} auto benchmark requests")
+    print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}")
+    print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}")
+    return dataset

python/sglang/benchmark/datasets/common.py

@@ -0,0 +1,101 @@
+import random
+from abc import ABC, abstractmethod
+from argparse import Namespace
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+ASSISTANT_SUFFIX = "Assistant:"
+SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered"
+SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json"
+MOONCAKE_DATASET_URL = {
+    "mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl",
+    "conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl",
+    "synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl",
+    "toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl",
+}
+
+
+@dataclass
+class DatasetRow:
+    prompt: Any
+    prompt_len: int
+    output_len: int
+    text_prompt_len: Optional[int] = None
+    vision_prompt_len: Optional[int] = None
+    image_data: Optional[List[str]] = None
+    timestamp: Optional[float] = None
+    routing_key: Optional[str] = None
+    extra_request_body: Optional[Dict[str, Any]] = None  # Per-request API parameters
+
+    def __post_init__(self):
+        if self.text_prompt_len is None:
+            self.text_prompt_len = self.prompt_len
+        if self.vision_prompt_len is None:
+            self.vision_prompt_len = 0
+        if self.extra_request_body is None:
+            self.extra_request_body = {}
+
+
+@dataclass
+class BaseDataset(ABC):
+    @classmethod
+    @abstractmethod
+    def from_args(cls, args: Namespace) -> "BaseDataset": ...
+
+    @abstractmethod
+    def load(
+        self,
+        tokenizer: Any,
+        model_id: Optional[str] = None,
+    ) -> List[DatasetRow]: ...
+
+
+def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]:
+    # full_len=0 is valid for embedding benchmarks where no output tokens are generated
+    if full_len <= 0:
+        return [0] * num
+    return np.random.randint(
+        max(int(full_len * range_ratio), 1),
+        full_len + 1,
+        size=num,
+    ).tolist()
+
+
+@lru_cache(maxsize=1)
+def get_available_tokens(tokenizer):
+    """Get valid token ids from the tokenizer vocabulary."""
+    return [
+        token_id
+        for token_id in tokenizer.get_vocab().values()
+        if isinstance(token_id, int)
+    ]
+
+
+def gen_prompt(tokenizer, token_num):
+    """Generate a random prompt of specified token length using tokenizer vocabulary."""
+    all_available_tokens = get_available_tokens(tokenizer)
+    selected_tokens = random.choices(all_available_tokens, k=token_num)
+    return tokenizer.decode(selected_tokens)
+
+
+@lru_cache(maxsize=1)
+def get_available_multimodal_text_tokens(tokenizer, image_pad_id):
+    """Get valid token ids for synthetic multimodal text prompts."""
+    excluded_token_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+    if image_pad_id is not None:
+        excluded_token_ids.add(image_pad_id)
+    return [
+        token_id
+        for token_id in get_available_tokens(tokenizer)
+        if token_id not in excluded_token_ids
+    ]
+
+
+def gen_mm_prompt(tokenizer, image_pad_id, token_num):
+    """Generate a random prompt of specified token length using tokenizer vocabulary."""
+    all_available_tokens = get_available_multimodal_text_tokens(tokenizer, image_pad_id)
+    selected_tokens = random.choices(all_available_tokens, k=token_num)
+    return tokenizer.decode(selected_tokens)

python/sglang/benchmark/datasets/custom.py

@@ -0,0 +1,147 @@
+import json
+import os
+import random
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    ASSISTANT_SUFFIX,
+    BaseDataset,
+    DatasetRow,
+)
+from sglang.benchmark.utils import remove_suffix
+
+
+@dataclass
+class CustomDataset(BaseDataset):
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+    context_len: Optional[int]
+    prompt_suffix: str
+    apply_chat_template: bool
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "CustomDataset":
+        assert not getattr(args, "tokenize_prompt", False)
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+            context_len=args.sharegpt_context_len,
+            prompt_suffix=args.prompt_suffix,
+            apply_chat_template=args.apply_chat_template,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_custom_requests(
+            dataset_path=self.dataset_path,
+            num_requests=self.num_requests,
+            tokenizer=tokenizer,
+            fixed_output_len=self.fixed_output_len,
+            context_len=self.context_len,
+            prompt_suffix=self.prompt_suffix,
+            apply_chat_template=self.apply_chat_template,
+        )
+
+
+def sample_custom_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: PreTrainedTokenizerBase,
+    fixed_output_len: Optional[int] = None,
+    context_len: Optional[int] = None,
+    prompt_suffix: Optional[str] = "",
+    apply_chat_template=False,
+) -> List[DatasetRow]:
+    """
+    Sample requests from a custom JSONL dataset: supports 'content'/'value' as conversation keys.
+    """
+    if fixed_output_len is not None and fixed_output_len < 4:
+        raise ValueError("output_len too small")
+
+    # Load the dataset
+    dataset = []
+    if not os.path.isfile(dataset_path):
+        raise FileNotFoundError(f"Dataset not found at {dataset_path}")
+
+    with open(dataset_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:  # skip empty lines
+                try:
+                    dataset.append(json.loads(line))
+                except json.JSONDecodeError:
+                    continue  # skip lines with JSON errors
+
+    # Filter out the conversations with less than 2 turns.
+    processed_dataset = []
+    for data in dataset:
+        convs = data.get("conversations", data.get("conversation", []))
+        if len(convs) >= 2:
+            user_turn = convs[0].get("content", convs[0].get("value", ""))
+            assist_turn = convs[1].get("content", convs[1].get("value", ""))
+            processed_dataset.append((user_turn, assist_turn))
+    dataset = processed_dataset
+    random.shuffle(dataset)
+
+    # Filter out sequences that are too long or too short
+    filtered_dataset: List[DatasetRow] = []
+
+    for i in range(len(dataset)):
+        if len(filtered_dataset) == num_requests:
+            break
+
+        # Tokenize the prompts and completions.
+        prompt = dataset[i][0]
+
+        if prompt_suffix:
+            prompt = (
+                remove_suffix(prompt, ASSISTANT_SUFFIX)
+                + prompt_suffix
+                + ASSISTANT_SUFFIX
+            )
+
+        if apply_chat_template:
+            prompt = tokenizer.apply_chat_template(
+                [{"role": "user", "content": prompt}],
+                add_generation_prompt=True,
+                tokenize=False,
+                return_dict=False,
+            )
+            if tokenizer.bos_token:
+                prompt = prompt.replace(tokenizer.bos_token, "")
+
+        prompt_token_ids = tokenizer.encode(prompt)
+        completion = dataset[i][1]
+        completion_token_ids = tokenizer.encode(completion)
+        prompt_len = len(prompt_token_ids)
+        output_len = (
+            len(completion_token_ids) if fixed_output_len is None else fixed_output_len
+        )
+
+        if prompt_len < 2 or output_len < 2:
+            # Prune too short sequences.
+            continue
+
+        if context_len and prompt_len + output_len > context_len:
+            # Prune too long sequences.
+            continue
+
+        filtered_dataset.append(
+            DatasetRow(
+                prompt=prompt,
+                prompt_len=prompt_len,
+                output_len=output_len,
+            )
+        )
+
+    print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
+    print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
+    return filtered_dataset

python/sglang/benchmark/datasets/generated_shared_prefix.py

@@ -0,0 +1,328 @@
+import math
+import pickle
+import random
+import uuid
+from argparse import Namespace
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+from tqdm.asyncio import tqdm
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    BaseDataset,
+    DatasetRow,
+    compute_random_lens,
+    gen_prompt,
+)
+
+
+def _zipf_group_probs(num_groups: int, alpha: float) -> np.ndarray:
+    """Rank-based Zipf probability vector with rank starting at 1.
+
+    weight(rank)      = 1 / rank ** alpha       (rank in 1..num_groups)
+    probability(rank) = weight(rank) / sum_over_all_ranks(weight)
+
+    The returned array has length num_groups; element i corresponds to
+    group index i (rank i + 1), so group 0 is the hottest.
+    """
+    if num_groups <= 0:
+        raise ValueError(f"num_groups must be > 0, got {num_groups}")
+    ranks = np.arange(1, num_groups + 1, dtype=np.float64)
+    weights = 1.0 / (ranks**alpha)
+    return weights / weights.sum()
+
+
+@dataclass
+class GeneratedSharedPrefixDataset(BaseDataset):
+    num_groups: int
+    prompts_per_group: int
+    system_prompt_len: int
+    question_len: int
+    output_len: int
+    range_ratio: float
+    seed: int
+    fast_prepare: bool
+    send_routing_key: bool
+    num_turns: int
+    ordered: bool
+    group_distribution: str = "uniform"
+    zipf_alpha: Optional[float] = None
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset":
+        assert not getattr(args, "tokenize_prompt", False)
+        group_distribution = getattr(args, "gsp_group_distribution", "uniform")
+        zipf_alpha = getattr(args, "gsp_zipf_alpha", None)
+
+        # Defensive validation for in-process callers that construct a
+        # Namespace by hand and bypass the argparse boundary in
+        # serving.py. The CLI hook enforces the same rules first.
+        if group_distribution not in ("uniform", "zipf"):
+            raise ValueError(
+                f"--gsp-group-distribution must be 'uniform' or 'zipf', "
+                f"got {group_distribution!r}"
+            )
+        if group_distribution == "zipf":
+            if zipf_alpha is None:
+                raise ValueError(
+                    "--gsp-group-distribution=zipf requires --gsp-zipf-alpha "
+                    "(a finite float > 0)"
+                )
+            if not math.isfinite(zipf_alpha) or zipf_alpha <= 0:
+                raise ValueError(
+                    f"--gsp-zipf-alpha must be a finite float > 0, got {zipf_alpha!r}"
+                )
+        elif zipf_alpha is not None:
+            raise ValueError(
+                "--gsp-zipf-alpha is only meaningful with "
+                "--gsp-group-distribution=zipf; remove --gsp-zipf-alpha "
+                "or set --gsp-group-distribution=zipf"
+            )
+
+        return cls(
+            num_groups=args.gsp_num_groups,
+            prompts_per_group=args.gsp_prompts_per_group,
+            system_prompt_len=args.gsp_system_prompt_len,
+            question_len=args.gsp_question_len,
+            output_len=args.gsp_output_len,
+            range_ratio=getattr(args, "gsp_range_ratio", 1.0),
+            seed=args.seed,
+            fast_prepare=getattr(args, "gsp_fast_prepare", False),
+            send_routing_key=getattr(args, "gsp_send_routing_key", False),
+            num_turns=getattr(args, "gsp_num_turns", 1),
+            ordered=getattr(args, "gsp_ordered", False),
+            group_distribution=group_distribution,
+            zipf_alpha=zipf_alpha,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_generated_shared_prefix_requests(
+            num_groups=self.num_groups,
+            prompts_per_group=self.prompts_per_group,
+            system_prompt_len=self.system_prompt_len,
+            question_len=self.question_len,
+            output_len=self.output_len,
+            range_ratio=self.range_ratio,
+            tokenizer=tokenizer,
+            seed=self.seed,
+            send_routing_key=self.send_routing_key,
+            num_turns=self.num_turns,
+            fast_prepare=self.fast_prepare,
+            ordered=self.ordered,
+            group_distribution=self.group_distribution,
+            zipf_alpha=self.zipf_alpha,
+        )
+
+
+def get_gen_prefix_cache_path(
+    seed: int,
+    num_groups: int,
+    prompts_per_group: int,
+    system_prompt_len: int,
+    question_len: int,
+    output_len: int,
+    tokenizer,
+    group_distribution: str = "uniform",
+    zipf_alpha: Optional[float] = None,
+):
+    """Create cache directory under ~/.cache/sglang/benchmark.
+
+    The uniform-mode filename is preserved exactly as before so existing
+    on-disk caches remain valid. Non-default sampling modes get an extra
+    suffix encoding the parameters that affect the cached payload.
+    """
+    cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"
+
+    suffix = ""
+    if group_distribution != "uniform":
+        suffix = f"_{group_distribution}_{zipf_alpha}"
+
+    cache_key = (
+        f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_"
+        f"{system_prompt_len}_{question_len}_{output_len}{suffix}_"
+        f"{tokenizer.__class__.__name__}.pkl"
+    )
+    return cache_dir / cache_key
+
+
+def sample_generated_shared_prefix_requests(
+    num_groups: int,
+    prompts_per_group: int,
+    system_prompt_len: int,
+    question_len: int,
+    output_len: int,
+    range_ratio: float,
+    tokenizer: PreTrainedTokenizerBase,
+    seed: int,
+    send_routing_key: bool = False,
+    num_turns: int = 1,
+    fast_prepare: bool = False,
+    ordered: bool = False,
+    group_distribution: str = "uniform",
+    zipf_alpha: Optional[float] = None,
+) -> List[DatasetRow]:
+    """Generate benchmark requests with shared system prompts using random tokens and caching.
+
+    When group_distribution is "uniform" (default), each group receives exactly
+    prompts_per_group requests; behavior matches the legacy generator.
+
+    When group_distribution is "zipf", each request's group is sampled by rank
+    with probability 1/rank**zipf_alpha / sum_k(1/k**zipf_alpha); rank starts at
+    1 and group index 0 is the hottest. Sampling uses an isolated
+    numpy.random.default_rng(seed) so the shared question/system-prompt pool
+    stays byte-identical to uniform mode for the same seed and other args.
+    Zipf mode is cached on disk under a distinct key per (group_distribution,
+    zipf_alpha) value.
+    """
+    cache_path = get_gen_prefix_cache_path(
+        seed,
+        num_groups,
+        prompts_per_group,
+        system_prompt_len,
+        question_len,
+        output_len,
+        tokenizer,
+        group_distribution=group_distribution,
+        zipf_alpha=zipf_alpha,
+    )
+    # range_ratio != 1 / num_turns > 1 perturb the payload but are not in the
+    # cache key; send_routing_key embeds a per-run uuid + timestamp that is
+    # meaningless to cache. Bypass for these pre-existing reasons only.
+    should_cache = range_ratio == 1 and not send_routing_key and num_turns == 1
+
+    if should_cache and cache_path.exists():
+        print(f"\nLoading cached generated input data from {cache_path}")
+        with open(cache_path, "rb") as f:
+            return pickle.load(f)
+
+    if not should_cache:
+        print(f"\nCache bypassed ({range_ratio=}, {send_routing_key=}, {num_turns=})")
+
+    print(
+        f"\nGenerating new input data... "
+        f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=}, {group_distribution=}, {zipf_alpha=})"
+    )
+
+    run_random_str = uuid.uuid4().hex[:8]
+    run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+
+    system_prompt_lens = compute_random_lens(
+        full_len=system_prompt_len,
+        range_ratio=range_ratio,
+        num=num_groups,
+    )
+    question_lens = np.array(
+        compute_random_lens(
+            full_len=question_len,
+            range_ratio=range_ratio,
+            num=num_groups * prompts_per_group * num_turns,
+        )
+    ).reshape(num_groups, prompts_per_group, num_turns)
+    output_lens = np.array(
+        compute_random_lens(
+            full_len=output_len,
+            range_ratio=range_ratio,
+            num=num_groups * prompts_per_group,
+        )
+    ).reshape(num_groups, prompts_per_group)
+    del system_prompt_len, question_len, output_len
+
+    system_prompts = [
+        gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups)
+    ]
+
+    # shape: (num_groups, prompts_per_group, num_turns)
+    questions = [
+        [
+            [
+                gen_prompt(tokenizer, int(question_lens[g, p, t]))
+                for t in range(num_turns)
+            ]
+            for p in range(prompts_per_group)
+        ]
+        for g in range(num_groups)
+    ]
+
+    # Per-slot group assignment. Uniform mode is the identity assignment
+    # [0,0,...,1,1,...,N-1,N-1]; zipf mode samples from the rank distribution
+    # using an isolated RNG so the module-level random / numpy.random state
+    # that compute_random_lens / gen_prompt rely on is never perturbed -- this
+    # keeps the system-prompt and question pool byte-identical to uniform mode
+    # for the same seed and other args.
+    total_slots = num_groups * prompts_per_group
+    if group_distribution == "uniform":
+        assignment = np.repeat(np.arange(num_groups), prompts_per_group)
+    else:  # "zipf"
+        rng = np.random.default_rng(seed)
+        probs = _zipf_group_probs(num_groups, zipf_alpha)
+        assignment = rng.choice(num_groups, size=total_slots, replace=True, p=probs)
+
+    input_requests = []
+    total_input_tokens = 0
+    total_output_tokens = 0
+    for slot_idx, sampled_g in enumerate(
+        tqdm(assignment, desc="Generating shared-prefix prompts")
+    ):
+        # src_(g,p) walks the question pool in uniform-enumeration order, so
+        # per-slot question text is reproducibly identical across modes.
+        src_g, src_p = divmod(slot_idx, prompts_per_group)
+        sampled_g = int(sampled_g)
+
+        system_prompt = system_prompts[sampled_g]
+        routing_key = (
+            f"{run_random_str}_{run_start_timestamp}_{sampled_g}"
+            if send_routing_key
+            else None
+        )
+        turn_questions = questions[src_g][src_p]
+        turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[1:]
+        full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts
+        prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0]))
+        output_len_val = int(output_lens[src_g, src_p])
+
+        input_requests.append(
+            DatasetRow(
+                prompt=full_prompt,
+                prompt_len=prompt_len,
+                output_len=output_len_val,
+                routing_key=routing_key,
+            )
+        )
+        total_input_tokens += prompt_len
+        total_output_tokens += output_len_val
+
+    if not ordered:
+        random.shuffle(input_requests)
+
+    print(f"\nGenerated shared prefix dataset statistics:")
+    print(f"Number of groups: {num_groups}")
+    print(f"Prompts per group: {prompts_per_group}")
+    print(f"Number of turns: {num_turns}")
+    print(f"Group distribution: {group_distribution}")
+    if group_distribution == "zipf":
+        print(f"Zipf alpha: {zipf_alpha}")
+    print(f"Total prompts: {len(input_requests)}")
+    if not fast_prepare:
+        print(f"Total input tokens: {total_input_tokens}")
+        print(f"Total output tokens: {total_output_tokens}")
+        print(
+            f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
+        )
+        all_questions = [q for group in questions for conv in group for q in conv]
+        print(
+            f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n"
+        )
+
+    if should_cache:
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        print(f"Caching generated input data to {cache_path}")
+        with open(cache_path, "wb") as f:
+            pickle.dump(input_requests, f)
+
+    return input_requests

python/sglang/benchmark/datasets/image.py

@@ -0,0 +1,381 @@
+import io
+import warnings
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import numpy as np
+import pybase64
+from PIL import Image
+from transformers import AutoProcessor
+
+from sglang.benchmark.datasets.common import (
+    BaseDataset,
+    DatasetRow,
+    compute_random_lens,
+    gen_mm_prompt,
+)
+from sglang.benchmark.utils import get_processor
+
+
+@dataclass
+class ImageDataset(BaseDataset):
+    num_requests: int
+    image_count: int
+    input_len: int
+    output_len: int
+    range_ratio: float
+    image_content: str
+    image_format: str
+    image_resolution: str
+    backend: str
+    random_image_count: bool
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "ImageDataset":
+        return cls(
+            num_requests=args.num_prompts,
+            image_count=args.image_count,
+            input_len=args.random_input_len,
+            output_len=args.random_output_len,
+            range_ratio=args.random_range_ratio,
+            image_content=args.image_content,
+            image_format=args.image_format,
+            image_resolution=args.image_resolution,
+            backend=args.backend,
+            random_image_count=args.random_image_count,
+        )
+
+    def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
+        processor = get_processor(model_id)
+        return sample_image_requests(
+            num_requests=self.num_requests,
+            image_count=self.image_count,
+            input_len=self.input_len,
+            output_len=self.output_len,
+            range_ratio=self.range_ratio,
+            processor=processor,
+            image_content=self.image_content,
+            image_format=self.image_format,
+            image_resolution=self.image_resolution,
+            backend=self.backend,
+            random_image_count=self.random_image_count,
+        )
+
+
+def parse_image_resolution(image_resolution: str) -> Tuple[int, int]:
+    """Parse image resolution into (width, height).
+
+    Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format
+    (e.g., '1080x1920' means height=1080, width=1920).
+    """
+    resolution_to_size = {
+        "4k": (3840, 2160),
+        "1080p": (1920, 1080),
+        "720p": (1280, 720),
+        "360p": (640, 360),
+    }
+    if image_resolution in resolution_to_size:
+        return resolution_to_size[image_resolution]
+
+    res = image_resolution.strip().lower()
+    if "x" in res:
+        parts = res.split("x")
+        if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
+            height = int(parts[0])
+            width = int(parts[1])
+            if height > 0 and width > 0:
+                return (width, height)
+
+    raise ValueError(
+        f"Unsupported image resolution: {image_resolution}. "
+        "Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)."
+    )
+
+
+def parse_random_image_resolution(
+    image_resolution: str,
+) -> Optional[Tuple[Tuple[int, int], Tuple[int, int]]]:
+    """Parse ``random:<min_h>x<min_w>-<max_h>x<max_w>`` image bounds.
+
+    Returns ``None`` for fixed resolutions. The returned dimensions are
+    ``(width, height)`` pairs, matching :func:`parse_image_resolution`.
+    """
+
+    prefix = "random:"
+    if not image_resolution.strip().lower().startswith(prefix):
+        return None
+
+    bounds = image_resolution.strip()[len(prefix) :].split("-", maxsplit=1)
+    if len(bounds) != 2:
+        raise ValueError(
+            "Random image resolution must be 'random:<min_h>x<min_w>-"
+            "<max_h>x<max_w>', for example 'random:256x256-1024x1024'."
+        )
+
+    min_width, min_height = parse_image_resolution(bounds[0])
+    max_width, max_height = parse_image_resolution(bounds[1])
+    if min_width > max_width or min_height > max_height:
+        raise ValueError("Random image resolution minimum cannot exceed maximum.")
+    return (min_width, min_height), (max_width, max_height)
+
+
+def create_mm_data_row(
+    text_prompt, images: list, images_base64, output_len, processor, backend
+):
+    try:
+        if type(processor).__name__ == "Phi4MMProcessor":
+            # <|endoftext10|> is the image token used in the phi-4-multimodal model.
+            content_items = text_prompt.replace("image 1", "|endoftext10|")
+        else:
+            content_items = [
+                {"type": "image", "image": {"url": image_base64}}
+                for image_base64 in images_base64
+            ]
+            content_items.append({"type": "text", "text": text_prompt})
+        prompt_str = processor.apply_chat_template(
+            [{"role": "user", "content": content_items}],
+            add_generation_prompt=True,
+            tokenize=False,
+        )
+    except Exception as e:
+        # Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL)
+        print(f"Error applying chat template: {e}, fallback to <image> tag")
+        # Some tokenizers do not support list content; fall back to a placeholder in the text
+        if type(processor).__name__ == "MiniCPMOProcessor":
+            prompt_str = f"(<image>./</image>){text_prompt}"
+        else:
+            prompt_str = f"<image>{text_prompt}"
+
+    # Calculate total tokens (text + vision)
+    if type(processor).__name__ == "KimiK25Processor":
+        medias = [{"type": "image", "image": img} for img in images]
+        prompt_len = processor(
+            text=prompt_str,
+            medias=medias,
+            return_tensors="pt",
+        )["input_ids"].numel()
+    elif type(processor).__name__ == "VLChatProcessor":
+        prompt_len = processor(
+            prompt=prompt_str,
+            images=images,
+            force_batchify=False,
+        )["input_ids"].numel()
+    elif type(processor).__name__ == "DeepseekVLV2Processor":
+        result = processor(
+            conversations=prompt_str,
+            images=images,
+            inference_mode=True,
+        )
+        prompt_len = result.input_ids.numel()
+    else:
+        prompt_len = processor(
+            text=[prompt_str],
+            images=images,
+            padding=False,
+            return_tensors="pt",
+        )["input_ids"].numel()
+
+    # Calculate text-only tokens
+    try:
+        # Create text-only version of the prompt
+        text_only_prompt = processor.apply_chat_template(
+            [{"role": "user", "content": text_prompt}],
+            add_generation_prompt=True,
+            tokenize=False,
+        )
+        text_prompt_len = processor(
+            text=[text_only_prompt],
+            padding=False,
+            return_tensors="pt",
+        )["input_ids"].numel()
+    except Exception:
+        # Fallback: just tokenize the text prompt directly
+        tokenizer_to_use = (
+            processor.tokenizer if hasattr(processor, "tokenizer") else processor
+        )
+        text_prompt_len = len(tokenizer_to_use.encode(text_prompt))
+
+    # Vision tokens = total tokens - text tokens
+    vision_prompt_len = prompt_len - text_prompt_len
+
+    supported_backends = [
+        "sglang",
+        "sglang-native",
+        "sglang-oai-chat",
+        "vllm-chat",
+    ]
+    if backend not in supported_backends:
+        raise ValueError(
+            f"Image dataset only supports backends: {supported_backends}, "
+            f"got '{backend}'."
+        )
+
+    # OpenAI chat handlers apply the chat template and receive images separately, so
+    # send the raw text. /generate does not apply a chat template, so it needs
+    # prompt_str, which contains the multimodal processor's image placeholders.
+    use_raw_prompt = backend in ("sglang-oai-chat", "vllm-chat")
+
+    return DatasetRow(
+        prompt=text_prompt if use_raw_prompt else prompt_str,
+        prompt_len=prompt_len,
+        output_len=output_len,
+        text_prompt_len=text_prompt_len,
+        vision_prompt_len=vision_prompt_len,
+        image_data=images_base64,
+    )
+
+
+def sample_image_requests(
+    num_requests: int,
+    image_count: int,
+    input_len: int,
+    output_len: int,
+    range_ratio: float,
+    processor: AutoProcessor,
+    image_content: str,
+    image_format: str,
+    image_resolution: str,
+    backend: str,
+    random_image_count: bool = False,
+) -> List[DatasetRow]:
+    """Generate requests with images.
+
+    - If ``random_image_count`` is True, each request includes a random number of images between 1 and ``image_count``.
+    - If ``random_image_count`` is False, each request includes exactly ``image_count`` images.
+    - Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p
+      (1280x720), 360p (640x360), custom ``heightxwidth`` (e.g.,
+      1080x1920), or ``random:<min_h>x<min_w>-<max_h>x<max_w>``.
+    - Text lengths follow the 'random' dataset sampling rule. ``prompt_len``
+      only counts text tokens and excludes image data.
+    """
+
+    random_resolution_bounds = parse_random_image_resolution(image_resolution)
+    if random_resolution_bounds is None:
+        width, height = parse_image_resolution(image_resolution)
+        min_width = max_width = width
+        min_height = max_height = height
+    else:
+        (min_width, min_height), (max_width, max_height) = random_resolution_bounds
+
+    # Determine image counts for each request
+    if random_image_count:
+        # Random number of images per request
+        image_counts = np.random.randint(1, image_count + 1, size=num_requests)
+        total_images = np.sum(image_counts)
+    else:
+        # Fixed number of images per request
+        image_counts = np.full(num_requests, image_count)
+        total_images = image_count * num_requests
+
+    # Check for potentially problematic combinations and warn user
+    if max_width * max_height >= 1920 * 1080 and total_images >= 100:
+        warnings.warn(
+            f"High resolution (up to {max_width}x{max_height}) with {total_images} total images "
+            f"may take a long time. Consider reducing resolution or image count.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    # Sample text lengths
+    input_lens = compute_random_lens(
+        full_len=input_len,
+        range_ratio=range_ratio,
+        num=num_requests,
+    )
+    output_lens = compute_random_lens(
+        full_len=output_len,
+        range_ratio=range_ratio,
+        num=num_requests,
+    )
+
+    def _gen_random_image_data_uri() -> Tuple[Image.Image, str, int, Tuple[int, int]]:
+        if random_resolution_bounds is None:
+            width, height = min_width, min_height
+        else:
+            width = np.random.randint(min_width, max_width + 1)
+            height = np.random.randint(min_height, max_height + 1)
+        if image_content == "blank":
+            # Generate blank white image
+            arr = np.full((height, width, 3), 255, dtype=np.uint8)
+        else:
+            # Generate random colored image
+            arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8)
+        img = Image.fromarray(arr)
+        buf = io.BytesIO()
+        img.save(buf, format=image_format, quality=85)
+        encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8")
+        image_data = f"data:image/{image_format};base64,{encoded}"
+        image_bytes = len(image_data.encode("utf-8"))
+        return img, image_data, image_bytes, (width, height)
+
+    dataset: List[DatasetRow] = []
+    total_image_bytes = 0
+    all_image_sizes: list[Tuple[int, int]] = []
+    for i in range(num_requests):
+        # Get the number of images for this request
+        request_image_count = int(image_counts[i])
+
+        # Generate text prompt
+        text_prompt = gen_mm_prompt(
+            processor.tokenizer if hasattr(processor, "tokenizer") else processor,
+            processor.image_token_id if hasattr(processor, "image_token_id") else None,
+            int(input_lens[i]),
+        )
+
+        # Generate image list
+        images, images_base64, images_bytes, image_sizes = zip(
+            *[_gen_random_image_data_uri() for _ in range(request_image_count)]
+        )
+        total_image_bytes += sum(images_bytes)
+        all_image_sizes.extend(image_sizes)
+
+        data_row = create_mm_data_row(
+            text_prompt,
+            list(images),
+            list(images_base64),
+            int(output_lens[i]),
+            processor,
+            backend,
+        )
+        dataset.append(data_row)
+
+    # Print statistics
+    print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}")
+    print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}")
+    print(f"#Total images: {total_images}")
+
+    if random_image_count:
+        print(
+            f"#Images per request: min={np.min(image_counts)}, max={np.max(image_counts)}, mean={np.mean(image_counts):.2f}"
+        )
+    else:
+        print(f"#Images per request: {image_count} (fixed)")
+
+    if random_resolution_bounds is not None:
+        widths, heights = zip(*all_image_sizes)
+        print(
+            "#Image resolution: "
+            f"min={min(widths)}x{min(heights)}, "
+            f"max={max(widths)}x{max(heights)}, "
+            f"mean={np.mean(widths):.1f}x{np.mean(heights):.1f}"
+        )
+
+    # Detailed token breakdown (derived from dataset + input_lens)
+    text_prompt_lens = np.array([r.text_prompt_len for r in dataset])
+    vision_prompt_lens = np.array([r.vision_prompt_len for r in dataset])
+    text_prompt_overheads = text_prompt_lens - input_lens
+    stat_fields = [
+        ("Raw text prompt tokens (without overhead)", input_lens),
+        ("Text prompt tokens (with chat template)", text_prompt_lens),
+        ("Text prompt overhead", text_prompt_overheads),
+        ("Vision tokens", vision_prompt_lens),
+    ]
+    print("\n=== Token Breakdown (per request avg / total) ===")
+    for label, vals in stat_fields:
+        print(f"  {label}: avg={np.mean(vals):.1f}, total={np.sum(vals)}")
+
+    print(
+        f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request"
+    )
+    return dataset

python/sglang/benchmark/datasets/longbench_v2.py

@@ -0,0 +1,104 @@
+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
+
+LONGBENCH_V2_REPO_ID = "THUDM/LongBench-v2"
+LONGBENCH_V2_DEFAULT_OUTPUT_LEN = 10  # answer letter + short explanation
+
+
+def _format_prompt(example: dict) -> str:
+    return (
+        f"{example['context']}\n\n"
+        f"Question: {example['question']}\n"
+        f"A. {example['choice_A']}\n"
+        f"B. {example['choice_B']}\n"
+        f"C. {example['choice_C']}\n"
+        f"D. {example['choice_D']}\n"
+        f"Answer:"
+    )
+
+
+@dataclass
+class LongBenchV2Dataset(BaseDataset):
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+    context_len: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "LongBenchV2Dataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+            context_len=args.sharegpt_context_len,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_longbench_v2_requests(
+            dataset_path=self.dataset_path,
+            num_requests=self.num_requests,
+            tokenizer=tokenizer,
+            fixed_output_len=self.fixed_output_len,
+            context_len=self.context_len,
+        )
+
+
+def sample_longbench_v2_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: PreTrainedTokenizerBase,
+    fixed_output_len: Optional[int] = None,
+    context_len: Optional[int] = None,
+) -> List[DatasetRow]:
+    output_len = (
+        fixed_output_len
+        if fixed_output_len is not None
+        else LONGBENCH_V2_DEFAULT_OUTPUT_LEN
+    )
+
+    # Load dataset
+    if dataset_path:
+        # Local file (parquet or JSON lines)
+        import pandas as pd
+
+        if dataset_path.endswith(".parquet"):
+            df = pd.read_parquet(dataset_path)
+            examples = df.to_dict(orient="records")
+        else:
+            import json
+
+            with open(dataset_path) as f:
+                examples = [json.loads(line) for line in f if line.strip()]
+    else:
+        from datasets import load_dataset
+
+        ds = load_dataset(LONGBENCH_V2_REPO_ID, split="train")
+        examples = list(ds)
+
+    random.shuffle(examples)
+
+    rows: List[DatasetRow] = []
+    for example in examples:
+        if len(rows) >= num_requests:
+            break
+
+        prompt = _format_prompt(example)
+        prompt_ids = tokenizer(prompt).input_ids
+        prompt_len = len(prompt_ids)
+
+        if context_len is not None and prompt_len + output_len > context_len:
+            continue
+
+        rows.append(
+            DatasetRow(prompt=prompt, prompt_len=prompt_len, output_len=output_len)
+        )
+
+    return rows

python/sglang/benchmark/datasets/mmmu.py

@@ -0,0 +1,124 @@
+import io
+import random
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import pybase64
+from datasets import load_dataset
+from transformers import AutoProcessor, AutoTokenizer
+
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+from sglang.benchmark.datasets.image import create_mm_data_row
+from sglang.benchmark.utils import get_processor
+
+
+@dataclass
+class MMMUDataset(BaseDataset):
+    num_requests: int
+    backend: str
+    fixed_output_len: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "MMMUDataset":
+        return cls(
+            num_requests=args.num_prompts,
+            backend=args.backend,
+            fixed_output_len=args.random_output_len,
+        )
+
+    def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
+        processor = get_processor(model_id)
+        return sample_mmmu_requests(
+            num_requests=self.num_requests,
+            processor=processor,
+            backend=self.backend,
+            fixed_output_len=self.fixed_output_len,
+        )
+
+
+def sample_mmmu_requests(
+    num_requests: int,
+    processor: AutoProcessor | AutoTokenizer,
+    backend: str = "sglang",
+    fixed_output_len: Optional[int] = None,
+    random_sample: bool = True,
+) -> List[DatasetRow]:
+    """
+    Sample requests from the MMMU dataset using HuggingFace datasets.
+
+    Args:
+        num_requests: Number of requests to sample.
+        fixed_output_len: If provided, use this fixed output length for all requests.
+        random_sample: Whether to randomly sample or take the first N.
+
+    Returns:
+        List of tuples (prompt, prompt_token_len, output_token_len).
+    """
+    print("Loading MMMU dataset from HuggingFace...")
+
+    try:
+        print("Attempting to load MMMU Math dataset...")
+        mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test")
+        print(
+            f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples"
+        )
+    except Exception as e:
+        print(f"Failed to load MMMU Math dataset: {e}")
+        raise ValueError(f"Failed to load MMMU dataset: {e}")
+
+    # Sample from the dataset
+    if len(mmmu_dataset) > num_requests:
+        if random_sample:
+            # Random sample
+            indices = random.sample(range(len(mmmu_dataset)), num_requests)
+            sample_dataset = mmmu_dataset.select(indices)
+        else:
+            # Take first N
+            sample_dataset = mmmu_dataset.select(
+                range(min(num_requests, len(mmmu_dataset)))
+            )
+    else:
+        print(f"Dataset has less than {num_requests} examples, using all examples")
+        sample_dataset = mmmu_dataset
+
+    print(f"Selected {len(sample_dataset)} examples for benchmarking")
+
+    # Create prompts
+    filtered_dataset = []
+
+    for i, example in enumerate(sample_dataset):
+        try:
+            # Extract image_1
+            image = example.get("image_1")
+
+            if image is not None:
+                if hasattr(image, "save"):
+                    # Convert RGBA images to RGB before encoding
+                    if image.mode == "RGBA":
+                        image = image.convert("RGB")
+
+                    # Encode image to base64 (save as PNG to support palette/alpha modes)
+                    buffered = io.BytesIO()
+                    image.save(buffered, format="PNG")
+                    img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8")
+                    image_data = f"data:image/png;base64,{img_str}"
+                else:
+                    continue
+
+                # Extract the question
+                question = example.get("question")
+
+                # Construct the prompt
+                text_prompt = f"Question: {question}\n\nAnswer: "
+                output_len = fixed_output_len if fixed_output_len is not None else 256
+                data_row = create_mm_data_row(
+                    text_prompt, [image], [image_data], output_len, processor, backend
+                )
+                filtered_dataset.append(data_row)
+
+        except Exception as e:
+            print(f"Error processing example {i}: {e}")
+
+    print(f"\nCreated {len(filtered_dataset)} MMMU prompts")
+    return filtered_dataset

python/sglang/benchmark/datasets/mooncake.py

@@ -0,0 +1,123 @@
+import asyncio
+import json
+import os
+import time
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import AsyncGenerator, Dict, List
+
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    MOONCAKE_DATASET_URL,
+    BaseDataset,
+    DatasetRow,
+)
+from sglang.benchmark.utils import download_and_cache_file
+
+
+@dataclass
+class MooncakeDataset(BaseDataset):
+    dataset_path: str
+    mooncake_workload: str
+    num_requests: int
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "MooncakeDataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            mooncake_workload=args.mooncake_workload,
+            num_requests=args.num_prompts,
+        )
+
+    def load(self, tokenizer=None, model_id=None) -> List[Dict]:
+        if not self.dataset_path:
+            local_path = os.path.join("/tmp", self.mooncake_workload + "_trace.jsonl")
+        else:
+            local_path = self.dataset_path
+
+        if not os.path.exists(local_path):
+            download_and_cache_file(
+                MOONCAKE_DATASET_URL[self.mooncake_workload], local_path
+            )
+
+        with open(local_path, "r") as f:
+            all_requests_data = [json.loads(line) for line in f if line.strip()]
+
+        return all_requests_data[: self.num_requests]
+
+
+async def get_mooncake_request_over_time(
+    input_requests: List[Dict],
+    tokenizer: PreTrainedTokenizerBase,
+    slowdown_factor: float,
+    num_rounds: int,
+) -> AsyncGenerator[DatasetRow, None]:
+    """
+    An async generator that yields requests based on the timestamps in the Mooncake trace file,
+    with support for multi-round sessions.
+    """
+    if not input_requests:
+        return
+
+    input_requests.sort(key=lambda r: r["timestamp"])
+
+    start_time = time.perf_counter()
+    trace_start_time_ms = input_requests[0]["timestamp"]
+
+    for record in input_requests:
+        # Calculate when this entire session should start
+        relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0
+        target_arrival_time_s = relative_arrival_time_s * slowdown_factor
+
+        current_elapsed_time_s = time.perf_counter() - start_time
+        sleep_duration_s = target_arrival_time_s - current_elapsed_time_s
+        if sleep_duration_s > 0:
+            await asyncio.sleep(sleep_duration_s)
+
+        # Once the session starts, generate all rounds for it as a burst
+        # This simulates a user engaging in a multi-turn conversation
+
+        # Base user query constructed from hash_ids
+        user_query_base = ""
+        hash_ids = record.get("hash_ids", [])
+        for hash_id in hash_ids:
+            user_query_base += f"{hash_id}" + " ".join(
+                ["hi"] * 128
+            )  # Shorter for multi-round
+        user_query_base += "Tell me a story based on this context."
+
+        output_len_per_round = record.get("output_length", 256)
+        chat_history = []
+
+        for i in range(num_rounds):
+            # Add user query for the current round
+            chat_history.append(
+                {"role": "user", "content": f"Round {i + 1}: {user_query_base}"}
+            )
+
+            # Form the full prompt from history
+            try:
+                full_prompt_text = tokenizer.apply_chat_template(
+                    chat_history,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    return_dict=False,
+                )
+            except Exception:
+                full_prompt_text = "\n".join(
+                    [f"{msg['role']}: {msg['content']}" for msg in chat_history]
+                )
+
+            prompt_len = len(tokenizer.encode(full_prompt_text))
+
+            yield DatasetRow(
+                prompt=full_prompt_text,
+                prompt_len=prompt_len,
+                output_len=output_len_per_round,
+            )
+
+            # Add a placeholder assistant response for the next round's context
+            # We use a placeholder because we don't know the real response
+            placeholder_response = " ".join(["story"] * output_len_per_round)
+            chat_history.append({"role": "assistant", "content": placeholder_response})

python/sglang/benchmark/datasets/openai_dataset.py

@@ -0,0 +1,113 @@
+import json
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+
+
+@dataclass
+class OpenAIDataset(BaseDataset):
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "OpenAIDataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_openai_requests(
+            dataset_path=self.dataset_path,
+            num_requests=self.num_requests,
+            tokenizer=tokenizer,
+            fixed_output_len=self.fixed_output_len,
+        )
+
+
+def sample_openai_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: PreTrainedTokenizerBase,
+    fixed_output_len: Optional[int] = None,
+) -> List[DatasetRow]:
+    """
+    Load OpenAI-compatible chat completion requests from a JSONL file.
+
+    Each line should be a JSON object with:
+    - "messages": list of {"role": str, "content": str}
+    - "max_tokens": int (used as output_len if fixed_output_len not set)
+    - "tools": optional list of tool definitions
+    - "temperature": optional temperature value
+    - "top_p": optional top_p value
+    - Other OpenAI API parameters are also extracted and passed through
+    """
+    dataset = []
+    with open(dataset_path, "r") as f:
+        for line in f:
+            if num_requests > 0 and len(dataset) >= num_requests:
+                break
+            if line.strip():
+                try:
+                    dataset.append(json.loads(line))
+                except json.JSONDecodeError:
+                    # Skip invalid JSON lines
+                    continue
+
+    # Fields that should NOT be passed through extra_request_body
+    # These are either handled separately or are metadata
+    # max_tokens is excluded because it's handled via output_len -> max_completion_tokens
+    # max_completion_tokens is also excluded to avoid conflicts
+    EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"}
+
+    filtered_dataset: List[DatasetRow] = []
+    for data in dataset:
+        messages = data.get("messages", [])
+        if not messages:
+            continue
+
+        # Use max_tokens from the request, or fall back to fixed_output_len
+        output_len = fixed_output_len or data.get("max_tokens", 256)
+
+        # Extract extra request body parameters (tools, temperature, top_p, etc.)
+        extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS}
+
+        # Calculate prompt length by applying chat template
+        # This includes the messages but not the tools
+        prompt_len = len(
+            tokenizer.apply_chat_template(
+                messages, tokenize=True, add_generation_prompt=True
+            )
+        )
+
+        # If tools are present, we need to add their token count
+        # Tools are sent as part of the request and count toward input tokens
+        if "tools" in extra_body:
+            # Encode tools as JSON string to estimate token count
+            tools_str = json.dumps(extra_body["tools"])
+            tools_tokens = len(tokenizer.encode(tools_str))
+            prompt_len += tools_tokens
+
+        # Pass messages list directly - the serving benchmark handles List[Dict] prompts
+        filtered_dataset.append(
+            DatasetRow(
+                prompt=messages,
+                prompt_len=prompt_len,
+                output_len=output_len,
+                extra_request_body=extra_body,  # Store per-request parameters
+            )
+        )
+
+    print(f"Loaded {len(filtered_dataset)} OpenAI-format requests")
+    print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
+    print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
+    return filtered_dataset

python/sglang/benchmark/datasets/random.py

@@ -0,0 +1,167 @@
+import json
+import random
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    SHAREGPT_FILENAME,
+    SHAREGPT_REPO_ID,
+    BaseDataset,
+    DatasetRow,
+    compute_random_lens,
+)
+from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json
+
+
+@dataclass
+class RandomDataset(BaseDataset):
+    input_len: int
+    output_len: int
+    num_requests: int
+    range_ratio: float
+    dataset_path: str
+    return_text: bool
+    random_sample: bool
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "RandomDataset":
+        return cls(
+            input_len=args.random_input_len,
+            output_len=args.random_output_len,
+            num_requests=args.num_prompts,
+            range_ratio=args.random_range_ratio,
+            dataset_path=args.dataset_path,
+            return_text=not getattr(args, "tokenize_prompt", False),
+            random_sample=(args.dataset_name == "random"),
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_random_requests(
+            input_len=self.input_len,
+            output_len=self.output_len,
+            num_prompts=self.num_requests,
+            range_ratio=self.range_ratio,
+            tokenizer=tokenizer,
+            dataset_path=self.dataset_path,
+            random_sample=self.random_sample,
+            return_text=self.return_text,
+        )
+
+
+def sample_random_requests(
+    input_len: int,
+    output_len: int,
+    num_prompts: int,
+    range_ratio: float,
+    tokenizer: PreTrainedTokenizerBase,
+    dataset_path: str,
+    random_sample: bool = True,
+    return_text: bool = True,
+) -> List[DatasetRow]:
+    input_lens = compute_random_lens(
+        full_len=input_len,
+        range_ratio=range_ratio,
+        num=num_prompts,
+    )
+    output_lens = compute_random_lens(
+        full_len=output_len,
+        range_ratio=range_ratio,
+        num=num_prompts,
+    )
+
+    if return_text:
+        # Need to truncate input_len as server encode will add special token.
+        num_special_tokens = int(tokenizer.num_special_tokens_to_add())
+        for i in range(num_prompts):
+            input_lens[i] = max(1, input_lens[i] - num_special_tokens)
+
+    if random_sample:
+        # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens
+
+        # Download sharegpt if necessary
+        if not is_file_valid_json(dataset_path):
+            dataset_path = download_and_cache_hf_file(
+                repo_id=SHAREGPT_REPO_ID,
+                filename=SHAREGPT_FILENAME,
+            )
+
+        # Load the dataset.
+        with open(dataset_path) as f:
+            dataset = json.load(f)
+        # Filter out the conversations with less than 2 turns.
+        dataset = [
+            data
+            for data in dataset
+            if len(data.get("conversations", data.get("conversation", []))) >= 2
+        ]
+        # Only keep the first two turns of each conversation.
+        dataset = [
+            (
+                data.get("conversations", data.get("conversation", []))[0]["value"],
+                data.get("conversations", data.get("conversation", []))[1]["value"],
+            )
+            for data in dataset
+        ]
+        # Shuffle the dataset.
+        random.shuffle(dataset)
+
+        # Filter out sequences that are too long or too short
+        input_requests: List[DatasetRow] = []
+        for data in dataset:
+            i = len(input_requests)
+            if i == num_prompts:
+                break
+
+            # Tokenize the prompts and completions.
+            prompt = data[0]
+            prompt_token_ids = tokenizer.encode(prompt)
+            prompt_len = len(prompt_token_ids)
+
+            # Skip empty prompt
+            if prompt_len == 0:
+                continue
+
+            if prompt_len > input_lens[i]:
+                input_ids = prompt_token_ids[: input_lens[i]]
+            else:
+                ratio = (input_lens[i] + prompt_len - 1) // prompt_len
+                input_ids = (prompt_token_ids * ratio)[: input_lens[i]]
+            input_content = input_ids
+            if return_text:
+                input_content = tokenizer.decode(input_content)
+            input_requests.append(
+                DatasetRow(
+                    prompt=input_content,
+                    prompt_len=input_lens[i],
+                    output_len=output_lens[i],
+                )
+            )
+    else:
+        # Sample token ids from random integers. This can cause some NaN issues.
+        offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
+        input_requests = []
+        for i in range(num_prompts):
+            # Use int() to convert numpy.int64 to native Python int for JSON serialization
+            input_content = [
+                int((offsets[i] + i + j) % tokenizer.vocab_size)
+                for j in range(input_lens[i])
+            ]
+            if return_text:
+                input_content = tokenizer.decode(input_content)
+            input_requests.append(
+                DatasetRow(
+                    prompt=input_content,
+                    prompt_len=input_lens[i],
+                    output_len=output_lens[i],
+                )
+            )
+
+    print(f"#Input tokens: {np.sum(input_lens)}")
+    print(f"#Output tokens: {np.sum(output_lens)}")
+    return input_requests

python/sglang/benchmark/datasets/sharegpt.py

@@ -0,0 +1,151 @@
+import json
+import random
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    ASSISTANT_SUFFIX,
+    SHAREGPT_FILENAME,
+    SHAREGPT_REPO_ID,
+    BaseDataset,
+    DatasetRow,
+)
+from sglang.benchmark.utils import (
+    download_and_cache_hf_file,
+    is_file_valid_json,
+    remove_suffix,
+)
+
+
+@dataclass
+class ShareGPTDataset(BaseDataset):
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+    context_len: Optional[int]
+    prompt_suffix: str
+    apply_chat_template: bool
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "ShareGPTDataset":
+        assert not getattr(args, "tokenize_prompt", False)
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+            context_len=args.sharegpt_context_len,
+            prompt_suffix=args.prompt_suffix,
+            apply_chat_template=args.apply_chat_template,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        return sample_sharegpt_requests(
+            dataset_path=self.dataset_path,
+            num_requests=self.num_requests,
+            tokenizer=tokenizer,
+            fixed_output_len=self.fixed_output_len,
+            context_len=self.context_len,
+            prompt_suffix=self.prompt_suffix,
+            apply_chat_template=self.apply_chat_template,
+        )
+
+
+def sample_sharegpt_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: PreTrainedTokenizerBase,
+    fixed_output_len: Optional[int] = None,
+    context_len: Optional[int] = None,
+    prompt_suffix: Optional[str] = "",
+    apply_chat_template=False,
+) -> List[DatasetRow]:
+    if fixed_output_len is not None and fixed_output_len < 4:
+        raise ValueError("output_len too small")
+
+    # Download sharegpt if necessary
+    if not is_file_valid_json(dataset_path) and dataset_path == "":
+        dataset_path = download_and_cache_hf_file(
+            repo_id=SHAREGPT_REPO_ID,
+            filename=SHAREGPT_FILENAME,
+        )
+
+    # Load the dataset.
+    with open(dataset_path) as f:
+        dataset = json.load(f)
+
+    # Filter out the conversations with less than 2 turns.
+    dataset = [
+        data
+        for data in dataset
+        if len(data.get("conversations", data.get("conversation", []))) >= 2
+    ]
+    # Only keep the first two turns of each conversation.
+    dataset = [
+        (
+            data.get("conversations", data.get("conversation", []))[0]["value"],
+            data.get("conversations", data.get("conversation", []))[1]["value"],
+        )
+        for data in dataset
+    ]
+
+    # Shuffle the dataset.
+    random.shuffle(dataset)
+
+    # Filter out sequences that are too long or too short
+    filtered_dataset: List[DatasetRow] = []
+    for i in range(len(dataset)):
+        if len(filtered_dataset) == num_requests:
+            break
+
+        # Tokenize the prompts and completions.
+        prompt = dataset[i][0]
+        if prompt_suffix:
+            prompt = (
+                remove_suffix(prompt, ASSISTANT_SUFFIX)
+                + prompt_suffix
+                + ASSISTANT_SUFFIX
+            )
+
+        if apply_chat_template:
+            prompt = tokenizer.apply_chat_template(
+                [{"role": "user", "content": prompt}],
+                add_generation_prompt=True,
+                tokenize=False,
+                return_dict=False,
+            )
+            if tokenizer.bos_token:
+                prompt = prompt.replace(tokenizer.bos_token, "")
+
+        prompt_token_ids = tokenizer.encode(prompt)
+        completion = dataset[i][1]
+        completion_token_ids = tokenizer.encode(completion)
+        prompt_len = len(prompt_token_ids)
+        output_len = (
+            len(completion_token_ids) if fixed_output_len is None else fixed_output_len
+        )
+
+        if prompt_len < 2 or output_len < 2:
+            # Prune too short sequences.
+            continue
+
+        if context_len and prompt_len + output_len > context_len:
+            # Prune too long sequences.
+            continue
+
+        filtered_dataset.append(
+            DatasetRow(
+                prompt=prompt,
+                prompt_len=prompt_len,
+                output_len=output_len,
+            )
+        )
+
+    print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}")
+    print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}")
+    return filtered_dataset

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