[Refactor] Benchmark Phase 1: extract utils and datasets from bench_serving (#19077)

Co-authored-by: Xuchun Shang <107600043+xucsh@users.noreply.github.com>

python/sglang/benchmark/datasets/__init__.py

@@ -0,0 +1,156 @@
+import json
+import os
+
+from sglang.benchmark.datasets.common import (
+    ASSISTANT_SUFFIX,
+    MOONCAKE_DATASET_URL,
+    SHAREGPT_FILENAME,
+    SHAREGPT_REPO_ID,
+    DatasetRow,
+    compute_random_lens,
+    gen_mm_prompt,
+    gen_prompt,
+    get_available_tokens,
+)
+from sglang.benchmark.datasets.custom import sample_custom_requests
+from sglang.benchmark.datasets.generated_shared_prefix import (
+    get_gen_prefix_cache_path,
+    sample_generated_shared_prefix_requests,
+)
+from sglang.benchmark.datasets.image import (
+    create_mm_data_row,
+    parse_image_resolution,
+    sample_image_requests,
+)
+from sglang.benchmark.datasets.mmmu import sample_mmmu_requests
+from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time
+from sglang.benchmark.datasets.openai_dataset import sample_openai_requests
+from sglang.benchmark.datasets.random import sample_random_requests
+from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
+from sglang.benchmark.utils import download_and_cache_file, get_processor
+
+
+def get_dataset(args, tokenizer, model_id=None):
+    tokenize_prompt = getattr(args, "tokenize_prompt", False)
+    if args.dataset_name == "sharegpt":
+        assert not tokenize_prompt
+        input_requests = sample_sharegpt_requests(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            tokenizer=tokenizer,
+            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,
+        )
+    elif args.dataset_name.startswith("random"):
+        input_requests = sample_random_requests(
+            input_len=args.random_input_len,
+            output_len=args.random_output_len,
+            num_prompts=args.num_prompts,
+            range_ratio=args.random_range_ratio,
+            tokenizer=tokenizer,
+            dataset_path=args.dataset_path,
+            random_sample=args.dataset_name == "random",
+            return_text=not tokenize_prompt,
+        )
+    elif args.dataset_name == "image":
+        processor = get_processor(model_id)
+        input_requests = sample_image_requests(
+            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,
+            processor=processor,
+            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,
+        )
+    elif args.dataset_name == "generated-shared-prefix":
+        assert not tokenize_prompt
+        input_requests = sample_generated_shared_prefix_requests(
+            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),
+            tokenizer=tokenizer,
+            args=args,
+        )
+    elif args.dataset_name == "mmmu":
+        processor = get_processor(model_id)
+        input_requests = sample_mmmu_requests(
+            num_requests=args.num_prompts,
+            processor=processor,
+            backend=args.backend,
+            fixed_output_len=args.random_output_len,
+            random_sample=True,
+        )
+    elif args.dataset_name == "mooncake":
+        # For mooncake, we don't generate the prompts here.
+        # We just load the raw trace data. The async generator will handle the rest.
+        if not args.dataset_path:
+            local_path = os.path.join("/tmp", args.mooncake_workload + "_trace.jsonl")
+        else:
+            local_path = args.dataset_path
+
+        if not os.path.exists(local_path):
+            download_and_cache_file(
+                MOONCAKE_DATASET_URL[args.mooncake_workload], local_path
+            )
+
+        with open(local_path, "r") as f:
+            all_requests_data = [json.loads(line) for line in f if line.strip()]
+
+        # Limit the number of requests based on --num-prompts
+        input_requests = all_requests_data[: args.num_prompts]
+    elif args.dataset_name == "custom":
+        assert not tokenize_prompt
+        input_requests = sample_custom_requests(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            tokenizer=tokenizer,
+            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,
+        )
+    elif args.dataset_name == "openai":
+        input_requests = sample_openai_requests(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            tokenizer=tokenizer,
+            fixed_output_len=args.sharegpt_output_len,
+        )
+    else:
+        raise ValueError(f"Unknown dataset: {args.dataset_name}")
+    return input_requests
+
+
+__all__ = [
+    "ASSISTANT_SUFFIX",
+    "MOONCAKE_DATASET_URL",
+    "SHAREGPT_FILENAME",
+    "SHAREGPT_REPO_ID",
+    "DatasetRow",
+    "compute_random_lens",
+    "create_mm_data_row",
+    "gen_mm_prompt",
+    "gen_prompt",
+    "get_available_tokens",
+    "get_dataset",
+    "get_gen_prefix_cache_path",
+    "get_mooncake_request_over_time",
+    "parse_image_resolution",
+    "sample_custom_requests",
+    "sample_generated_shared_prefix_requests",
+    "sample_image_requests",
+    "sample_mmmu_requests",
+    "sample_openai_requests",
+    "sample_random_requests",
+    "sample_sharegpt_requests",
+]

python/sglang/benchmark/datasets/common.py

@@ -0,0 +1,67 @@
+import random
+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 = {}
+
+
+def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]:
+    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 all available token ids from the tokenizer vocabulary."""
+    return list(tokenizer.get_vocab().values())
+
+
+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)
+
+
+def gen_mm_prompt(tokenizer, image_pad_id, token_num):
+    """Generate a random prompt of specified token length using tokenizer vocabulary."""
+    all_available_tokens = list(tokenizer.get_vocab().values())
+    if image_pad_id:
+        all_available_tokens.remove(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,106 @@
+import json
+import os
+import random
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import ASSISTANT_SUFFIX, DatasetRow
+from sglang.benchmark.utils import remove_suffix
+
+
+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,163 @@
+import argparse
+import pickle
+import random
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import List
+
+import numpy as np
+from tqdm.asyncio import tqdm
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import DatasetRow, compute_random_lens, gen_prompt
+
+
+def get_gen_prefix_cache_path(args, tokenizer):
+    """Create cache directory under ~/.cache/sglang/benchmark"""
+    cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"
+
+    # Create a unique cache filename based on the generation parameters
+    cache_key = (
+        f"gen_shared_prefix_{args.seed}_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_"
+        f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_"
+        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,
+    args: argparse.Namespace,
+) -> List[DatasetRow]:
+    """Generate benchmark requests with shared system prompts using random tokens and caching."""
+    send_routing_key = getattr(args, "gsp_send_routing_key", False)
+    num_turns = getattr(args, "gsp_num_turns", 1)
+
+    cache_path = get_gen_prefix_cache_path(args, tokenizer)
+    should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1
+
+    # Try to load from cache first
+    if cache_path.exists() and should_cache:
+        print(f"\nLoading cached generated input data from {cache_path}")
+        with open(cache_path, "rb") as f:
+            return pickle.load(f)
+
+    print(
+        f"\nGenerating new input data... "
+        f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})"
+    )
+
+    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
+
+    # Generate system prompts for each group
+    system_prompts = [
+        gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups)
+    ]
+
+    # Generate questions: 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)
+    ]
+
+    # Combine system prompts with questions
+    input_requests = []
+    total_input_tokens = 0
+    total_output_tokens = 0
+
+    for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
+        system_prompt = system_prompts[group_idx]
+        routing_key = (
+            f"{run_random_str}_{run_start_timestamp}_{group_idx}"
+            if send_routing_key
+            else None
+        )
+        for prompt_idx in tqdm(
+            range(prompts_per_group), desc="Generating questions", leave=False
+        ):
+            turn_questions = questions[group_idx][prompt_idx]
+            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 getattr(args, "gsp_fast_prepare", False)
+                else len(tokenizer.encode(turn_prompts[0]))
+            )
+            output_len_val = int(output_lens[group_idx, prompt_idx])
+
+            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 getattr(args, "gsp_ordered", False):
+        random.shuffle(input_requests)
+
+    # Print statistics
+    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"Total prompts: {len(input_requests)}")
+    if not getattr(args, "gsp_fast_prepare", False):
+        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"
+        )
+
+    # Save to cache
+    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,239 @@
+import io
+import warnings
+from typing import List, Tuple
+
+import numpy as np
+import pybase64
+from PIL import Image
+from transformers import AutoProcessor
+
+from sglang.benchmark.datasets.common import (
+    DatasetRow,
+    compute_random_lens,
+    gen_mm_prompt,
+)
+
+
+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 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
+        prompt_str = f"<image>{text_prompt}"
+
+    # Calculate total tokens (text + vision)
+    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
+
+    use_raw_prompt = backend in [
+        "sglang",
+        "sglang-oai",
+        "sglang-oai-chat",
+        "vllm",
+        "vllm-chat",
+        "lmdeploy",
+        "lmdeploy-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),
+      or custom 'heightxwidth' (e.g., 1080x1920).
+    - Text lengths follow the 'random' dataset sampling rule. ``prompt_len``
+      only counts text tokens and excludes image data.
+    """
+
+    # Parse resolution (supports presets and 'heightxwidth')
+    width, height = parse_image_resolution(image_resolution)
+
+    # 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 width * height >= 1920 * 1080 and total_images >= 100:
+        warnings.warn(
+            f"High resolution ({width}x{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(
+        width: int = width, height: int = height
+    ) -> Tuple[Image.Image, str, int]:
+        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
+
+    dataset: List[DatasetRow] = []
+    total_image_bytes = 0
+    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,
+            processor.image_token_id if hasattr(processor, "image_token_id") else None,
+            int(input_lens[i]),
+        )
+
+        # Generate image list
+        images, images_base64, images_bytes = zip(
+            *[_gen_random_image_data_uri() for _ in range(request_image_count)]
+        )
+        total_image_bytes += sum(images_bytes)
+
+        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)")
+
+    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/mmmu.py

@@ -0,0 +1,97 @@
+import io
+import random
+from typing import List, Optional
+
+import pybase64
+from datasets import load_dataset
+from transformers import AutoProcessor, AutoTokenizer
+
+from sglang.benchmark.datasets.common import DatasetRow
+from sglang.benchmark.datasets.image import create_mm_data_row
+
+
+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,83 @@
+import asyncio
+import time
+from typing import AsyncGenerator, Dict, List
+
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import DatasetRow
+
+
+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,86 @@
+import json
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import DatasetRow
+
+
+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 - bench_serving 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,127 @@
+import json
+import random
+from typing import List
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import (
+    SHAREGPT_FILENAME,
+    SHAREGPT_REPO_ID,
+    DatasetRow,
+    compute_random_lens,
+)
+from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json
+
+
+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(0, 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,113 @@
+import json
+import random
+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,
+    DatasetRow,
+)
+from sglang.benchmark.utils import (
+    download_and_cache_hf_file,
+    is_file_valid_json,
+    remove_suffix,
+)
+
+
+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/utils.py

@@ -0,0 +1,159 @@
+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 != ""
+    )
+    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_processor
+
+        return get_processor(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 AutoProcessor.from_pretrained(
+        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}")