chore: import upstream snapshot with attribution

python/sglang/benchmark/bench_utils.py

@@ -0,0 +1,23 @@
+"""Triton do_bench/do_bench_cudagraph compatible wrapper using flashinfer.testing.bench_gpu_time."""
+
+import numpy as np
+from flashinfer.testing import bench_gpu_time
+
+
+def run_bench(
+    fn,
+    use_cuda_graph: bool = True,
+    quantiles=(0.5, 0.2, 0.8),
+    warmup_ms: int = 25,
+    rep_ms: int = 100,
+):
+    """Returns (ms, min_ms, max_ms) or (median,) when quantiles=None."""
+    times = bench_gpu_time(
+        fn=fn,
+        use_cuda_graph=use_cuda_graph,
+        dry_run_time_ms=warmup_ms,
+        repeat_time_ms=rep_ms,
+    )
+    if quantiles is None:
+        return (float(np.median(times)),)
+    return tuple(float(np.percentile(times, q * 100)) for q in quantiles)

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

python/sglang/benchmark/dspark_sts_fit.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import argparse
+import glob
+import logging
+import math
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_sts import (
+    DSparkStsCalibration,
+)
+
+logger = logging.getLogger(__name__)
+
+_EPS_PROB = 1e-8
+
+
+def default_temperature_grid() -> torch.Tensor:
+    return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41)
+
+
+def expected_calibration_error(
+    *,
+    probs: torch.Tensor,
+    targets: torch.Tensor,
+    num_bins: int,
+) -> float:
+    probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB)
+    targets = targets.reshape(-1).to(torch.float64)
+    total = probs.numel()
+    if total == 0:
+        return float("nan")
+    bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1)
+    count = torch.zeros(num_bins, dtype=torch.float64)
+    pred_sum = torch.zeros(num_bins, dtype=torch.float64)
+    target_sum = torch.zeros(num_bins, dtype=torch.float64)
+    count.scatter_add_(0, bin_index, torch.ones_like(probs))
+    pred_sum.scatter_add_(0, bin_index, probs)
+    target_sum.scatter_add_(0, bin_index, targets)
+    denom = count.clamp_min(1.0)
+    bin_error = (pred_sum / denom - target_sum / denom).abs()
+    return float((bin_error * count).sum().item() / total)
+
+
+def fit_sts_temperatures(
+    *,
+    logits: torch.Tensor,
+    prefix_mask: torch.Tensor,
+    grid: torch.Tensor,
+    num_bins: int = 15,
+) -> dict[str, list[float]]:
+    logits = logits.to(torch.float64)
+    prefix_mask = prefix_mask.to(torch.float64)
+    num_samples, gamma = logits.shape
+    if num_samples == 0:
+        raise ValueError("fit_sts_temperatures requires at least one sample.")
+    grid_values = grid.to(torch.float64).tolist()
+
+    temperatures: list[float] = []
+    ece_before: list[float] = []
+    ece_after: list[float] = []
+
+    survival_at_one = torch.ones(num_samples, dtype=torch.float64)
+    survival_fitted = torch.ones(num_samples, dtype=torch.float64)
+    for position in range(gamma):
+        position_logits = logits[:, position]
+        position_target = prefix_mask[:, position]
+
+        survival_at_one = survival_at_one * torch.sigmoid(position_logits)
+        ece_before.append(
+            expected_calibration_error(
+                probs=survival_at_one,
+                targets=position_target,
+                num_bins=num_bins,
+            )
+        )
+
+        best_temperature = grid_values[0]
+        best_survival = survival_fitted * torch.sigmoid(
+            position_logits / best_temperature
+        )
+        best_ece = expected_calibration_error(
+            probs=best_survival, targets=position_target, num_bins=num_bins
+        )
+        for temperature in grid_values[1:]:
+            candidate_survival = survival_fitted * torch.sigmoid(
+                position_logits / temperature
+            )
+            candidate_ece = expected_calibration_error(
+                probs=candidate_survival,
+                targets=position_target,
+                num_bins=num_bins,
+            )
+            if candidate_ece < best_ece:
+                best_ece = candidate_ece
+                best_temperature = temperature
+                best_survival = candidate_survival
+
+        temperatures.append(float(best_temperature))
+        ece_after.append(float(best_ece))
+        survival_fitted = best_survival
+
+    return {
+        "temperatures": temperatures,
+        "ece_before": ece_before,
+        "ece_after": ece_after,
+    }
+
+
+def load_collected_shards(*, data_glob: str) -> tuple[torch.Tensor, torch.Tensor]:
+    shard_paths = sorted(glob.glob(data_glob))
+    if not shard_paths:
+        raise ValueError(f"No STS data shards matched {data_glob!r}.")
+
+    logits_shards: list[torch.Tensor] = []
+    prefix_mask_shards: list[torch.Tensor] = []
+    shard_gamma: Optional[int] = None
+    for shard_path in shard_paths:
+        shard = torch.load(shard_path, map_location="cpu")
+        shard_logits = shard["logits"]
+        shard_prefix_mask = shard["prefix_mask"]
+        if shard_logits.shape != shard_prefix_mask.shape:
+            raise ValueError(
+                f"Shard {shard_path!r} logits / prefix_mask shape mismatch: "
+                f"{tuple(shard_logits.shape)} vs {tuple(shard_prefix_mask.shape)}."
+            )
+        if shard_gamma is None:
+            shard_gamma = int(shard_logits.shape[1])
+        elif int(shard_logits.shape[1]) != shard_gamma:
+            raise ValueError(
+                f"Shard {shard_path!r} gamma {int(shard_logits.shape[1])} disagrees "
+                f"with earlier shards' gamma {shard_gamma}."
+            )
+        logits_shards.append(shard_logits)
+        prefix_mask_shards.append(shard_prefix_mask)
+
+    return torch.cat(logits_shards, dim=0), torch.cat(prefix_mask_shards, dim=0)
+
+
+def fit(
+    *,
+    data_glob: str,
+    out: Path,
+    num_bins: int = 15,
+    gamma: Optional[int] = None,
+) -> None:
+    logits, prefix_mask = load_collected_shards(data_glob=data_glob)
+    resolved_gamma = int(logits.shape[1])
+    if gamma is not None and gamma != resolved_gamma:
+        raise ValueError(
+            f"Collected shards have gamma={resolved_gamma} but --gamma={gamma}."
+        )
+    num_samples = int(logits.shape[0])
+
+    result = fit_sts_temperatures(
+        logits=logits,
+        prefix_mask=prefix_mask,
+        grid=default_temperature_grid(),
+        num_bins=num_bins,
+    )
+    calibration = DSparkStsCalibration(
+        temperatures=result["temperatures"],
+        dataset=data_glob,
+        num_samples=num_samples,
+        ece_before=result["ece_before"],
+        ece_after=result["ece_after"],
+    )
+    out.write_text(calibration.to_json(), encoding="utf-8")
+
+    print(
+        f"Fit STS temperatures over {num_samples} samples (gamma={resolved_gamma}) "
+        f"-> {out}"
+    )
+    print("pos  temperature  ece_before  ece_after")
+    for position in range(resolved_gamma):
+        print(
+            f"{position:>3}  {result['temperatures'][position]:>11.4f}  "
+            f"{result['ece_before'][position]:>10.4f}  "
+            f"{result['ece_after'][position]:>9.4f}"
+        )
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser(
+        description="Fit DSpark Sequential Temperature Scaling (STS) calibration "
+        "temperatures from collected confidence shards."
+    )
+    parser.add_argument(
+        "--data-glob",
+        required=True,
+        help="Glob of collected .pt shards, each a dict with [n, gamma] "
+        "'logits' and 'prefix_mask' tensors.",
+    )
+    parser.add_argument(
+        "--out",
+        required=True,
+        type=Path,
+        help="Output STS calibration JSON path.",
+    )
+    parser.add_argument(
+        "--num-bins",
+        type=int,
+        default=15,
+        help="Number of equal-width ECE bins.",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=int,
+        default=None,
+        help="Optional gamma override to validate the shards against.",
+    )
+    args = parser.parse_args()
+
+    fit(
+        data_glob=args.data_glob,
+        out=args.out,
+        num_bins=args.num_bins,
+        gamma=args.gamma,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/benchmark/endpoint.py

@@ -0,0 +1,127 @@
+"""Connection target for HTTP benchmark scripts.
+
+Owns the launch-vs-connect decision in one place: a benchmark only needs a base
+URL, which comes either from a server we launch or one already running.
+"""
+
+import dataclasses
+import multiprocessing
+import os
+import time
+from typing import Callable, Optional
+
+import requests
+
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.network import resolve_base_url
+
+DEFAULT_TIMEOUT = 600
+
+# Field defaults of ServerArgs, used to detect when --host/--port were set
+# explicitly (and would be silently ignored in connect mode).
+_SERVER_ARGS_DEFAULTS = {f.name: f.default for f in dataclasses.fields(ServerArgs)}
+
+
+def server_is_up(base_url: str, timeout: float = DEFAULT_TIMEOUT) -> bool:
+    """Return True if a server answers /v1/models with 200 at base_url."""
+    try:
+        headers = {
+            "Content-Type": "application/json; charset=utf-8",
+        }
+        response = requests.get(
+            f"{base_url}/v1/models", headers=headers, timeout=timeout
+        )
+        return response.status_code == 200
+    except requests.RequestException:
+        return False
+
+
+def _launch_server_target(launch_server_func: Callable, server_args: ServerArgs):
+    try:
+        launch_server_func(server_args)
+    except Exception as e:
+        raise e
+    finally:
+        kill_process_tree(os.getpid(), include_parent=False)
+
+
+def launch_or_reuse_server(launch_server_func: Callable, server_args: ServerArgs):
+    base_url = resolve_base_url("", server_args.host, server_args.port)
+
+    # Reuse an already-running server instead of forking a second one onto the
+    # occupied port, where it would orphan, compete for the GPU, and OOM.
+    if server_is_up(base_url, timeout=5):
+        print(
+            f"WARNING: reusing the server already running at {base_url} "
+            f"(--model and server-launch args ignored). Pass --base-url to silence."
+        )
+        return None, base_url
+
+    proc = multiprocessing.Process(
+        target=_launch_server_target,
+        args=(
+            launch_server_func,
+            server_args,
+        ),
+    )
+    proc.start()
+
+    start_time = time.time()
+    while time.time() - start_time < DEFAULT_TIMEOUT:
+        # Fail fast if the server died during startup (e.g. OOM).
+        if not proc.is_alive():
+            raise RuntimeError(
+                f"Server process exited during startup (exit code "
+                f"{proc.exitcode}); see the traceback above for the cause."
+            )
+        if server_is_up(base_url):
+            return proc, base_url
+        time.sleep(10)
+
+    # Timed out: kill the half-started server so it does not linger as an orphan.
+    kill_process_tree(proc.pid)
+    raise TimeoutError("Server failed to start within the timeout period.")
+
+
+@dataclasses.dataclass
+class BenchEndpoint:
+    """A base URL plus the lifecycle of any server we launched to back it.
+    ``close()`` tears down a launched server; for a connected one it is a no-op.
+    """
+
+    base_url: str
+    _proc: Optional[multiprocessing.Process] = None
+
+    def close(self) -> None:
+        if self._proc is not None:
+            kill_process_tree(self._proc.pid)
+            self._proc = None
+
+
+def acquire_endpoint(
+    server_args: ServerArgs,
+    base_url: str = "",
+    launch_server_func: Callable = launch_server,
+) -> BenchEndpoint:
+    """Resolve the benchmark target -- the single launch-vs-connect decision.
+
+    base_url given: connect to it (server_args is ignored). base_url empty:
+    launch a server from server_args. Caller must close() the result.
+    """
+    if base_url:
+        ignored = [
+            f"--{name}"
+            for name in ("host", "port")
+            if getattr(server_args, name) != _SERVER_ARGS_DEFAULTS[name]
+        ]
+        if ignored:
+            print(
+                f"WARNING: --base-url is set, so {' / '.join(ignored)} (and other "
+                f"launch args) are ignored; benchmarking the server at {base_url}."
+            )
+        return BenchEndpoint(base_url=base_url)
+
+    proc, url = launch_or_reuse_server(launch_server_func, server_args)
+    return BenchEndpoint(base_url=url, _proc=proc)

python/sglang/benchmark/offline_throughput.py

@@ -0,0 +1,586 @@
+"""
+Benchmark the throughput in the offline mode.
+It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as serving.py).
+
+# Usage
+## Sharegpt dataset with default args
+python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10
+
+## Random dataset with default args
+python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024
+
+## Random dataset with profiling args
+SGLANG_TORCH_PROFILER_DIR=/tmp python -m sglang.benchmark.offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 128 --random-output 128 --num-prompts 4 --max-running-requests 4 --profile-steps 3 --profile --profile-activities "CPU" "XPU"
+"""
+
+import argparse
+import asyncio
+import dataclasses
+import inspect
+import json
+import logging
+import os
+import random
+import time
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+from sglang.benchmark.datasets import DatasetRow, get_dataset
+from sglang.benchmark.datasets.random import sample_random_requests
+from sglang.benchmark.utils import get_tokenizer, set_ulimit
+from sglang.lang.backend.runtime_endpoint import Runtime
+from sglang.srt.entrypoints.engine import Engine
+from sglang.srt.server_args import ServerArgs
+
+
+@dataclasses.dataclass
+class BenchArgs:
+    backend: str = "engine"
+    result_filename: str = ""
+    dataset_name: str = "sharegpt"
+    dataset_path: str = ""
+    num_prompts: int = 1000
+    sharegpt_output_len: Optional[int] = None
+    sharegpt_context_len: Optional[int] = None
+    random_input_len: int = 1024
+    random_output_len: int = 1024
+    random_range_ratio: float = 0.0
+    gsp_num_groups: int = 64
+    gsp_prompts_per_group: int = 16
+    gsp_system_prompt_len: int = 2048
+    gsp_question_len: int = 128
+    gsp_output_len: int = 256
+    seed: int = 42
+    disable_ignore_eos: bool = False
+    extra_request_body: Optional[str] = None
+    apply_chat_template: bool = False
+    profile: bool = False
+    profile_activities: Tuple[str] = ("CPU", "GPU")
+    profile_steps: Optional[int] = None
+    skip_warmup: bool = False
+    do_not_exit: bool = False
+    prompt_suffix: str = ""
+    return_logprob: bool = False
+    logprob_start_len: int = -1
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser):
+        parser.add_argument("--backend", type=str, default=BenchArgs.backend)
+        parser.add_argument(
+            "--result-filename", type=str, default=BenchArgs.result_filename
+        )
+        parser.add_argument(
+            "--dataset-name",
+            type=str,
+            default="sharegpt",
+            choices=["sharegpt", "random", "generated-shared-prefix"],
+            help="Name of the dataset to benchmark on.",
+        )
+        parser.add_argument(
+            "--dataset-path", type=str, default="", help="Path to the dataset."
+        )
+        parser.add_argument(
+            "--num-prompts",
+            type=int,
+            default=BenchArgs.num_prompts,
+            help="Number of prompts to process. Default is 1000.",
+        )
+        parser.add_argument(
+            "--sharegpt-output-len",
+            type=int,
+            default=BenchArgs.sharegpt_output_len,
+            help="Output length for each request. Overrides the output length from the ShareGPT dataset.",
+        )
+        parser.add_argument(
+            "--sharegpt-context-len",
+            type=int,
+            default=BenchArgs.sharegpt_context_len,
+            help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.",
+        )
+        parser.add_argument(
+            "--random-input-len",
+            type=int,
+            default=BenchArgs.random_input_len,
+            help="Number of input tokens per request, used only for random dataset.",
+        )
+        parser.add_argument(
+            "--random-output-len",
+            type=int,
+            default=BenchArgs.random_output_len,
+            help="Number of output tokens per request, used only for random dataset.",
+        )
+        parser.add_argument(
+            "--random-range-ratio",
+            type=float,
+            default=BenchArgs.random_range_ratio,
+            help="Range of sampled ratio of input/output length, "
+            "used only for random dataset.",
+        )
+        parser.add_argument(
+            "--gsp-num-groups",
+            type=int,
+            default=BenchArgs.gsp_num_groups,
+            help="Number of groups with shared prefix, used"
+            "only for generate-shared-prefix",
+        )
+        parser.add_argument(
+            "--gsp-prompts-per-group",
+            type=int,
+            default=BenchArgs.gsp_prompts_per_group,
+            help="Number of prompts per group of shared prefix, used"
+            "only for generate-shared-prefix",
+        )
+        parser.add_argument(
+            "--gsp-system-prompt-len",
+            type=int,
+            default=BenchArgs.gsp_system_prompt_len,
+            help="System prompt length, used" "only for generate-shared-prefix",
+        )
+        parser.add_argument(
+            "--gsp-question-len",
+            type=int,
+            default=BenchArgs.gsp_question_len,
+            help="Question length, used" "only for generate-shared-prefix",
+        )
+        parser.add_argument(
+            "--gsp-output-len",
+            type=int,
+            default=BenchArgs.gsp_output_len,
+            help="Target length in tokens for outputs in generated-shared-prefix dataset",
+        )
+        parser.add_argument("--seed", type=int, default=42, help="The random seed.")
+        parser.add_argument(
+            "--disable-ignore-eos",
+            action="store_true",
+            help="Disable ignore EOS token",
+        )
+        parser.add_argument(
+            "--extra-request-body",
+            metavar='{"key1": "value1", "key2": "value2"}',
+            type=str,
+            default=BenchArgs.extra_request_body,
+            help="Append given JSON object to the request payload. You can use this to specify"
+            "additional generate params like sampling params.",
+        )
+        parser.add_argument(
+            "--apply-chat-template",
+            action="store_true",
+            help="Apply chat template",
+        )
+        parser.add_argument(
+            "--profile",
+            action="store_true",
+            help="Use Torch Profiler. The endpoint must be launched with "
+            "SGLANG_TORCH_PROFILER_DIR to enable profiler.",
+        )
+        parser.add_argument(
+            "--profile-activities",
+            type=str,
+            nargs="+",
+            default=["CPU", "GPU"],
+            choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"],
+            help="Profiler activities: CPU, GPU, XPU, CUDA_PROFILER. If CPU/GPU/XPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.",
+        )
+        parser.add_argument(
+            "--profile-steps",
+            type=int,
+            default=None,
+            help="Number of steps to profile. If not specified, profiles all steps.",
+        )
+        parser.add_argument(
+            "--skip-warmup",
+            action="store_true",
+            help="Skip the warmup batches.",
+        )
+        parser.add_argument(
+            "--do-not-exit",
+            action="store_true",
+            help="Do not exit the program. This is useful for nsys profile with --duration and --delay.",
+        )
+        parser.add_argument(
+            "--prompt-suffix",
+            type=str,
+            default="",
+            help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.",
+        )
+        parser.add_argument(
+            "--return-logprob",
+            action="store_true",
+            help="Enable returning log probabilities.",
+        )
+        parser.add_argument(
+            "--logprob-start-len",
+            type=int,
+            default=-1,
+            help="Start length for logprob. -1 means only return logprobs for output tokens (default). 0 means return logprobs for all tokens including input.",
+        )
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace):
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        return cls(**{attr: getattr(args, attr) for attr in attrs})
+
+
+def throughput_test_once(
+    backend_name: str,
+    backend,
+    reqs: List[DatasetRow],
+    ignore_eos: bool,
+    extra_request_body: Dict,
+    profile: bool,
+    profile_activities=None,
+    profile_steps=None,
+    return_logprob: bool = False,
+    logprob_start_len: int = -1,
+):
+    measurement_results = {
+        "backend": backend_name,
+        "successful_requests": len(reqs),
+        "total_latency": -1,
+        "total_input_tokens": sum(r.prompt_len for r in reqs),
+        "total_output_tokens": -1,
+        "request_throughput": -1,
+        "input_throughput": -1,
+        "output_throughput": -1,
+        "total_throughput": -1,
+    }
+
+    prompt = [r.prompt for r in reqs]
+    sampling_params = [
+        {
+            "temperature": 0,
+            "max_new_tokens": r.output_len,
+            "ignore_eos": ignore_eos,
+            **extra_request_body,
+        }
+        for r in reqs
+    ]
+
+    if profile:
+        assert (
+            "SGLANG_TORCH_PROFILER_DIR" in os.environ
+        ), "Please set SGLANG_TORCH_PROFILER_DIR."
+        os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True)
+        known_files = None
+        backend.start_profile(
+            num_steps=profile_steps,
+            activities=profile_activities,
+        )
+        if profile_steps:
+            dir = os.getenv("SGLANG_TORCH_PROFILER_DIR")
+            known_files = set(os.listdir(dir))
+
+    st = time.perf_counter()
+    gen_out = backend.generate(
+        prompt=prompt,
+        sampling_params=sampling_params,
+        return_logprob=return_logprob,
+        logprob_start_len=logprob_start_len,
+    )
+    latency = time.perf_counter() - st
+
+    if profile:
+        dir = os.getenv("SGLANG_TORCH_PROFILER_DIR")
+        if not profile_steps:
+            known_files = set(os.listdir(dir))
+        # With --profile-steps the scheduler auto-stops mid-run after N steps, so
+        # a second stop here raises "not in progress"; a run shorter than N steps
+        # never hit the target and still needs this explicit stop. Either way we
+        # must stop before monitor_trace_file, which loops forever waiting for a
+        # trace that would otherwise never be finalized.
+        try:
+            backend.stop_profile()
+        except RuntimeError:
+            pass
+        monitor_trace_file(known_files, dir)
+
+    if backend_name == "runtime":
+        gen_out = json.loads(gen_out)
+
+    server_info = backend.get_server_info()
+
+    measurement_results["total_latency"] = latency
+    measurement_results["total_output_tokens"] = sum(
+        o["meta_info"]["completion_tokens"] for o in gen_out
+    )
+    measurement_results["request_throughput"] = (
+        measurement_results["successful_requests"] / latency
+    )
+    measurement_results["input_throughput"] = (
+        measurement_results["total_input_tokens"] / latency
+    )
+    measurement_results["output_throughput"] = (
+        measurement_results["total_output_tokens"] / latency
+    )
+    measurement_results["total_throughput"] = (
+        measurement_results["total_input_tokens"]
+        + measurement_results["total_output_tokens"]
+    ) / latency
+
+    if inspect.isawaitable(server_info):
+        server_info = asyncio.run(server_info)
+
+    measurement_results["last_gen_throughput"] = server_info["internal_states"][0][
+        "last_gen_throughput"
+    ]
+
+    return measurement_results
+
+
+def monitor_trace_file(known_files, directory, interval=1):
+    print(f"Monitoring {directory} for new trace files...")
+
+    while True:
+        flag = False
+        time.sleep(interval)
+        current_files = set(os.listdir(directory))
+
+        new_files = current_files - known_files
+        for new_file in new_files:
+            new_file_path = os.path.join(directory, new_file)
+            print(f"New file detected: {new_file}")
+
+            previous_size = 0
+            while True:
+                try:
+                    current_size = os.path.getsize(new_file_path)
+                except FileNotFoundError:
+                    print(f"File {new_file} is no longer accessible.")
+                    break
+
+                if current_size > previous_size:
+                    previous_size = current_size
+                else:
+                    flag = True
+                    break
+
+                time.sleep(interval)
+        if flag:
+            break
+
+
+def _create_ray_engine_backend(server_args: ServerArgs):
+    """Create a RayEngine inside a Ray actor on a placement group.
+
+    RayEngine requires a placement group, so we launch it inside a Ray actor
+    and return a lightweight proxy that forwards calls via ray.get().
+    """
+    import ray
+    from ray.runtime_env import RuntimeEnv
+    from ray.util.placement_group import placement_group
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+    env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
+    if os.environ.get("HF_TOKEN"):
+        env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
+    if not ray.is_initialized():
+        ray.init(runtime_env=RuntimeEnv(env_vars=env_vars))
+
+    total_gpus = server_args.tp_size * server_args.pp_size
+    pg = placement_group([{"CPU": 1, "GPU": total_gpus}], strategy="STRICT_PACK")
+    ray.get(pg.ready())
+
+    @ray.remote
+    class _EngineActor:
+        def __init__(self, **kwargs):
+            from sglang.srt.ray.engine import RayEngine
+
+            self.engine = RayEngine(**kwargs)
+
+        def call(self, method, **kwargs):
+            return getattr(self.engine, method)(**kwargs)
+
+    actor = _EngineActor.options(
+        num_cpus=1,
+        num_gpus=0,
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_bundle_index=0,
+        ),
+    ).remote(**dataclasses.asdict(server_args))
+
+    class _Proxy:
+        """Forwards method calls to the remote RayEngine actor."""
+
+        def generate(self, **kwargs):
+            return ray.get(actor.call.remote("generate", **kwargs))
+
+        def get_server_info(self, **kwargs):
+            return ray.get(actor.call.remote("get_server_info", **kwargs))
+
+        def start_profile(self, **kwargs):
+            return ray.get(actor.call.remote("start_profile", **kwargs))
+
+        def stop_profile(self, **kwargs):
+            return ray.get(actor.call.remote("stop_profile", **kwargs))
+
+        def shutdown(self):
+            try:
+                ray.get(actor.call.remote("shutdown"), timeout=60)
+            except Exception:
+                pass
+            try:
+                ray.util.remove_placement_group(pg)
+            except Exception:
+                pass
+
+    return _Proxy()
+
+
+def throughput_test(
+    server_args: ServerArgs,
+    bench_args: BenchArgs,
+):
+    if bench_args.backend == "engine":
+        if server_args.use_ray:
+            backend = _create_ray_engine_backend(server_args)
+        else:
+            backend = Engine(**dataclasses.asdict(server_args))
+        if not backend:
+            raise ValueError("Please provide valid engine arguments")
+    elif bench_args.backend == "runtime":
+        backend = Runtime(**dataclasses.asdict(server_args))
+    else:
+        raise ValueError('Please set backend to either "engine" or "runtime"')
+
+    tokenizer_id = server_args.tokenizer_path or server_args.model_path
+    tokenizer = get_tokenizer(tokenizer_id)
+
+    # Set global environments
+    set_ulimit()
+    random.seed(bench_args.seed)
+    np.random.seed(bench_args.seed)
+
+    # Parse args
+    extra_request_body = {}
+    if bench_args.extra_request_body:
+        extra_request_body = json.loads(bench_args.extra_request_body)
+
+    # Read dataset
+    input_requests = get_dataset(bench_args, tokenizer)
+
+    warmup_requests = sample_random_requests(
+        input_len=256,
+        output_len=16,
+        num_prompts=min(bench_args.num_prompts, 16),
+        range_ratio=1.0,
+        tokenizer=tokenizer,
+        dataset_path=bench_args.dataset_path,
+    )
+
+    # Warm up
+    if not bench_args.skip_warmup:
+        logging.info("\nWarmup...")
+        throughput_test_once(
+            backend_name=bench_args.backend,
+            backend=backend,
+            reqs=warmup_requests,
+            ignore_eos=not bench_args.disable_ignore_eos,
+            extra_request_body=extra_request_body,
+            profile=False,
+            return_logprob=bench_args.return_logprob,
+            logprob_start_len=bench_args.logprob_start_len,
+        )
+        time.sleep(0.5)
+
+    logging.info("\nBenchmark...")
+    result = throughput_test_once(
+        backend_name=bench_args.backend,
+        backend=backend,
+        reqs=input_requests,
+        ignore_eos=not bench_args.disable_ignore_eos,
+        extra_request_body=extra_request_body,
+        profile=bench_args.profile,
+        profile_activities=bench_args.profile_activities,
+        profile_steps=bench_args.profile_steps,
+        return_logprob=bench_args.return_logprob,
+        logprob_start_len=bench_args.logprob_start_len,
+    )
+    backend.shutdown()
+
+    if bench_args.result_filename:
+        with open(bench_args.result_filename, "a") as fout:
+            fout.write(json.dumps(result) + "\n")
+
+    print(
+        "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=")
+    )
+    print("{:<40} {:<10}".format("Backend:", result["backend"]))
+    print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"]))
+    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"]))
+    print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"]))
+    print(
+        "{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"])
+    )
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Last generation throughput (tok/s):", result["last_gen_throughput"]
+        )
+    )
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Request throughput (req/s):", result["request_throughput"]
+        )
+    )
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Input token throughput (tok/s):", result["input_throughput"]
+        )
+    )
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Output token throughput (tok/s):", result["output_throughput"]
+        )
+    )
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Total token throughput (tok/s):", result["total_throughput"]
+        )
+    )
+    print("=" * 50)
+
+    return result
+
+
+def cli_main():
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    BenchArgs.add_cli_args(parser)
+    args = parser.parse_args()
+
+    # handling ModelScope model downloads
+    if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() in ("true", "1"):
+        if os.path.exists(args.model_path):
+            print(f"Using local model path: {args.model_path}")
+        else:
+            try:
+                from modelscope import snapshot_download
+
+                print(f"Using ModelScope to download model: {args.model_path}")
+
+                # download the model and replace args.model_path
+                args.model_path = snapshot_download(
+                    args.model_path,
+                )
+                print(f"Model downloaded to: {args.model_path}")
+            except Exception as e:
+                print(f"ModelScope download failed: {str(e)}")
+                raise e
+
+    server_args = ServerArgs.from_cli_args(args)
+    bench_args = BenchArgs.from_cli_args(args)
+
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    throughput_test(server_args, bench_args)
+
+    while bench_args.do_not_exit:
+        pass
+
+
+if __name__ == "__main__":
+    cli_main()

python/sglang/benchmark/utils.py

@@ -0,0 +1,158 @@
+import json
+import os
+import resource
+from json import JSONDecodeError
+from typing import Dict, List, Optional, Union
+
+import requests
+from tqdm.asyncio import tqdm
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+
+def remove_prefix(text: str, prefix: str) -> str:
+    return text[len(prefix) :] if text.startswith(prefix) else text
+
+
+def remove_suffix(text: str, suffix: str) -> str:
+    return text[: -len(suffix)] if text.endswith(suffix) else text
+
+
+def parse_custom_headers(header_list: List[str]) -> Dict[str, str]:
+    return {k: v for h in header_list for k, _, v in [h.partition("=")] if k and v}
+
+
+def get_model(pretrained_model_name_or_path: str) -> str:
+    if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true":
+        import huggingface_hub.constants
+        from modelscope import snapshot_download
+
+        model_path = snapshot_download(
+            model_id=pretrained_model_name_or_path,
+            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+            ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"],
+        )
+
+        return model_path
+    return pretrained_model_name_or_path
+
+
+def get_tokenizer(
+    pretrained_model_name_or_path: str,
+) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+    assert (
+        pretrained_model_name_or_path is not None
+        and pretrained_model_name_or_path != ""
+    )
+    if pretrained_model_name_or_path.endswith(
+        ".json"
+    ) or pretrained_model_name_or_path.endswith(".model"):
+        from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+
+        return get_tokenizer(pretrained_model_name_or_path)
+
+    if pretrained_model_name_or_path is not None and not os.path.exists(
+        pretrained_model_name_or_path
+    ):
+        pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
+    return AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path, trust_remote_code=True
+    )
+
+
+def get_processor(
+    pretrained_model_name_or_path: str,
+) -> AutoProcessor:
+    assert (
+        pretrained_model_name_or_path is not None
+        and pretrained_model_name_or_path != ""
+    )
+
+    from sglang.srt.utils.hf_transformers_utils import (
+        get_processor as _srt_get_processor,
+    )
+
+    if not pretrained_model_name_or_path.endswith(
+        (".json", ".model")
+    ) and not os.path.exists(pretrained_model_name_or_path):
+        pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
+    return _srt_get_processor(pretrained_model_name_or_path, trust_remote_code=True)
+
+
+def download_and_cache_hf_file(
+    repo_id: str,
+    filename: str,
+    repo_type: str = "dataset",
+):
+    """Download a file from Hugging Face and cache it locally."""
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type)
+
+
+def download_and_cache_file(url: str, filename: Optional[str] = None):
+    """Read and cache a file from a url."""
+    if filename is None:
+        filename = os.path.join("/tmp", url.split("/")[-1])
+
+    # Check if the cache file already exists
+    if is_file_valid_json(filename):
+        return filename
+
+    print(f"Downloading from {url} to {filename}")
+
+    # Stream the response to show the progress bar
+    response = requests.get(url, stream=True)
+    response.raise_for_status()  # Check for request errors
+
+    # Total size of the file in bytes
+    total_size = int(response.headers.get("content-length", 0))
+    chunk_size = 1024  # Download in chunks of 1KB
+
+    # Use tqdm to display the progress bar
+    with (
+        open(filename, "wb") as f,
+        tqdm(
+            desc=filename,
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar,
+    ):
+        for chunk in response.iter_content(chunk_size=chunk_size):
+            f.write(chunk)
+            bar.update(len(chunk))
+
+    return filename
+
+
+def is_file_valid_json(path):
+    if not os.path.isfile(path):
+        return False
+
+    # TODO can fuse into the real file open later
+    try:
+        with open(path) as f:
+            json.load(f)
+        return True
+    except JSONDecodeError as e:
+        print(
+            f"{path} exists but json loading fails ({e=}), thus treat as invalid file"
+        )
+        return False
+
+
+def set_ulimit(target_soft_limit=65535):
+    resource_type = resource.RLIMIT_NOFILE
+    current_soft, current_hard = resource.getrlimit(resource_type)
+
+    if current_soft < target_soft_limit:
+        try:
+            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
+        except ValueError as e:
+            print(f"Fail to set RLIMIT_NOFILE: {e}")