[Diffusion] Add mixed-resolution benchmark support (for #20762) (#20863)

Signed-off-by: Fengyuan Yu <15fengyuan@gmail.com>
Co-authored-by: Fengyuan Yu <15fengyuan@gmail.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/benchmarks/bench_offline_throughput.py

@@ -29,7 +29,7 @@ import dataclasses
 import json
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 from tqdm import tqdm
@@ -81,6 +81,8 @@ class BenchArgs:
     task_name: str = "unknown"
     num_prompts: int = 10
     batch_size: int = 1
+    random_request_config: str = None
+    random_request_seed: int = 42
 
     # Benchmark Execution
     skip_warmup: bool = False
@@ -151,6 +153,23 @@ class BenchArgs:
             help="Batch size per generation call (currently only bs=1 is supported)",
         )
 
+        parser.add_argument(
+            "--random-request-config",
+            type=str,
+            default=None,
+            help=(
+                "JSON string defining random request profiles. "
+                "Each profile may contain: width, height, num_inference_steps, etc. "
+                "The 'weight' field controls sampling probability (relative weight)."
+            ),
+        )
+        parser.add_argument(
+            "--random-request-seed",
+            type=int,
+            default=42,
+            help="Random seed for sampling request profiles (default: 42).",
+        )
+
         # Benchmark Execution
         parser.add_argument(
             "--skip-warmup", action="store_true", help="Skip warmup batch"
@@ -186,17 +205,22 @@ def generate_batch(
     engine: DiffGenerator,
     bench_args: BenchArgs,
     prompts: List[str],
-    user_sampling_params: Dict[str, Any],
+    user_sampling_params: List[Dict[str, Any]],
 ) -> BatchOutput:
     """Generate batch of images/videos synchronously."""
+    assert len(user_sampling_params) == len(prompts), (
+        f"user_sampling_params length ({len(user_sampling_params)}) must match "
+        f"prompts length ({len(prompts)})"
+    )
+
     output = BatchOutput()
     start_time = time.perf_counter()
 
     torch.cuda.reset_peak_memory_stats()
 
-    for prompt in prompts:
+    for prompt, params in zip(prompts, user_sampling_params):
         try:
-            sampling_params_kwargs = dict(user_sampling_params)
+            sampling_params_kwargs = dict(params)
             sampling_params_kwargs["prompt"] = prompt
             result = engine.generate(sampling_params_kwargs=sampling_params_kwargs)
 
@@ -227,6 +251,7 @@ def calculate_metrics(
     total_duration: float,
     resolution: Tuple[int, int, int],
     num_requests: int,
+    all_sampling_params: Optional[List[Dict[str, Any]]] = None,
 ) -> Dict[str, Any]:
     """Calculate generation-specific throughput metrics."""
     successful = [o for o in outputs if o.success]
@@ -235,8 +260,15 @@ def calculate_metrics(
     peak_memory = max((o.peak_memory_mb for o in outputs), default=0)
 
     width, height, frames = resolution
-    pixels_per_sample = width * height * frames
+    if all_sampling_params:
-    total_pixels = num_success * pixels_per_sample
+        total_pixels = sum(
+            p.get("width", width)
+            * p.get("height", height)
+            * p.get("num_frames", frames)
+            for p in all_sampling_params[:num_success]
+        )
+    else:
+        total_pixels = num_success * width * height * frames
 
     metrics = {
         "num_requests": num_requests,
@@ -272,6 +304,11 @@ def throughput_test(
 
     engine = initialize_engine(server_args)
 
+    if bench_args.random_request_config and bench_args.dataset != "random":
+        raise ValueError(
+            "--random-request-config can only be used with --dataset random"
+        )
+
     logger.info(f"Loading {bench_args.dataset} dataset...")
     if bench_args.dataset == "vbench":
         bench_args.task_name = engine.server_args.pipeline_config.task_type
@@ -281,7 +318,7 @@ def throughput_test(
     else:
         raise ValueError(f"Unknown dataset: {bench_args.dataset}")
 
-    sampling_params = {
+    _sampling_params = {
         "guidance_scale": bench_args.guidance_scale,
         "num_inference_steps": bench_args.num_inference_steps,
         "height": bench_args.height,
@@ -290,18 +327,29 @@ def throughput_test(
         "seed": bench_args.seed,
     }
     if bench_args.disable_safety_checker:
-        sampling_params["safety_checker"] = None
+        _sampling_params["safety_checker"] = None
+
+    total_count = min(bench_args.num_prompts, len(dataset))
+    all_prompts = [dataset[i].prompt for i in range(total_count)]
+
+    if bench_args.random_request_config:
+        all_sampling_params = []
+        for i in range(total_count):
+            params = dict(_sampling_params)
+            params.update(dataset.get_sampling_params(i))
+            all_sampling_params.append(params)
+    else:
+        all_sampling_params = [_sampling_params] * total_count
 
     if not bench_args.skip_warmup:
         logger.info("Running warmup batch...")
-        warmup_count = min(bench_args.batch_size, len(dataset))
+        warmup_count = min(bench_args.batch_size, total_count)
-        warmup_prompts = [dataset[i].prompt for i in range(warmup_count)]
+        warmup_prompts = all_prompts[:warmup_count]
-        generate_batch(engine, bench_args, warmup_prompts, sampling_params)
+        warmup_sampling_params = all_sampling_params[:warmup_count]
+        generate_batch(engine, bench_args, warmup_prompts, warmup_sampling_params)
 
     logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...")
     outputs: List[BatchOutput] = []
-    total_count = min(bench_args.num_prompts, len(dataset))
-    all_prompts = [dataset[i].prompt for i in range(total_count)]
 
     start_time = time.perf_counter()
 
@@ -315,9 +363,10 @@ def throughput_test(
     for batch_start in range(0, total_count, bench_args.batch_size):
         batch_end = min(batch_start + bench_args.batch_size, total_count)
         batch_prompts = all_prompts[batch_start:batch_end]
+        batch_sampling_params = all_sampling_params[batch_start:batch_end]
 
         batch_output = generate_batch(
-            engine, bench_args, batch_prompts, sampling_params
+            engine, bench_args, batch_prompts, batch_sampling_params
         )
         outputs.append(batch_output)
 
@@ -332,6 +381,7 @@ def throughput_test(
         total_duration,
         resolution=resolution,
         num_requests=total_count,
+        all_sampling_params=all_sampling_params,
     )
 
     display_results(

python/sglang/multimodal_gen/benchmarks/bench_serving.py

@@ -191,6 +191,8 @@ async def async_request_image_sglang(
 
         if input.width and input.height:
             payload["size"] = f"{input.width}x{input.height}"
+        if input.num_inference_steps:
+            payload["num_inference_steps"] = input.num_inference_steps
 
         # Merge extra parameters
         payload.update(input.extra_body)
@@ -299,6 +301,8 @@ async def async_request_video_sglang(
             payload["size"] = f"{input.width}x{input.height}"
         if input.num_frames:
             payload["num_frames"] = input.num_frames
+        if input.num_inference_steps:
+            payload["num_inference_steps"] = input.num_inference_steps
         if input.fps:
             payload["fps"] = input.fps
 
@@ -519,6 +523,11 @@ async def benchmark(args):
 
     setattr(args, "task_name", task_name)
 
+    if args.random_request_config and args.dataset != "random":
+        raise ValueError(
+            "--random-request-config can only be used with --dataset random"
+        )
+
     if args.dataset == "vbench":
         dataset = VBenchDataset(args, api_url, args.model)
     elif args.dataset == "random":
@@ -720,6 +729,25 @@ if __name__ == "__main__":
     )
     parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
     parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
+    parser.add_argument(
+        "--random-request-config",
+        type=str,
+        default=None,
+        help=(
+            "JSON string defining random request profiles. "
+            "Each profile may contain: width, height, num_inference_steps, etc. "
+            "The 'weight' field controls sampling probability (relative weight). "
+            "Example: "
+            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
+            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
+        ),
+    )
+    parser.add_argument(
+        "--random-request-seed",
+        type=int,
+        default=42,
+        help="Random seed for sampling request profiles (default: 42).",
+    )
     parser.add_argument(
         "--num-frames", type=int, default=None, help="Number of frames (for video)."
     )

python/sglang/multimodal_gen/benchmarks/datasets.py

@@ -1,6 +1,7 @@
 import glob
 import json
 import os
+import random
 import re
 import subprocess
 import uuid
@@ -286,16 +287,39 @@ class RandomDataset(BaseDataset):
         super().__init__(args, api_url, model)
         self.num_prompts = args.num_prompts or 100
 
+        self.random_request_config = args.random_request_config
+        if self.random_request_config:
+            self.random_request_config = json.loads(self.random_request_config)
+            weights = [p.pop("weight") for p in self.random_request_config]
+            seed = args.random_request_seed
+            rng = random.Random(seed)
+            self._sampled_requests = rng.choices(
+                self.random_request_config, weights=weights, k=self.num_prompts
+            )
+        else:
+            self._sampled_requests = None
+
+    def get_sampling_params(self, idx: int) -> dict:
+        """Return the per-request sampling profile dict, or empty dict if not mix-diffusion."""
+        if self._sampled_requests:
+            return self._sampled_requests[idx]
+        return {}
+
     def __len__(self) -> int:
         return self.num_prompts
 
     def __getitem__(self, idx: int) -> RequestFuncInput:
+        profile = self._sampled_requests[idx] if self._sampled_requests else {}
+
         return RequestFuncInput(
             prompt=f"Random prompt {idx} for benchmarking diffusion models",
             api_url=self.api_url,
             model=self.model,
-            width=self.args.width,
+            width=profile.get("width", self.args.width),
-            height=self.args.height,
+            height=profile.get("height", self.args.height),
-            num_frames=self.args.num_frames,
+            num_frames=profile.get("num_frames", self.args.num_frames),
-            fps=self.args.fps,
+            num_inference_steps=profile.get(
+                "num_inference_steps", self.args.num_inference_steps
+            ),
+            fps=profile.get("fps", self.args.fps),
         )