[Benchmark] DiT Model Benchmark under Mixed Workloads (#529)

Signed-off-by: asumaska <asumasaka@noreply.gitcode.com>
Co-authored-by: asumaska <asumasaka@noreply.gitcode.com>
This commit is contained in:
TsuKKIandasumaska authored and GitHub committed 2026-01-09 12:36:53 +08:00
1 parent 978e3ad7d4
commit f26219162b
2 files changed
+539 -15

No files matched your search

+117
View File
@@ -0,0 +1,117 @@
# Diffusion Serving Benchmark (Image/Video)
This folder contains an online-serving benchmark script for diffusion models.
It sends requests to a vLLM OpenAI-compatible endpoint and reports throughput,
latency percentiles, and optional SLO attainment.
The main entrypoint is:
- `benchmarks/diffusion/diffusion_benchmark_serving.py`
## 1. Quick Start
1. Start the server:
```bash
vllm serve Qwen/Qwen-Image --omni --port 8099
```
2. Run a minimal benchmark:
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset vbench \
--num-prompts 5
```
**Notes**
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
- If you run the server on another host or port, pass `--base-url` accordingly.
## 2. Supported Datasets
The benchmark supports three dataset modes via `--dataset`:
- `vbench`: Built-in prompt/data loader.
- `trace`: Heterogeneous request traces (each request can have different resolution/frames/steps).
- `random`: Synthetic prompts for quick smoke tests.
### VBench dataset
If you use i2v/i2i bench datasets and need auto-download support, you may need:
```bash
uv pip install gdown
```
### Trace dataset
Use `--dataset trace` to replay a trace file. The trace can specify per-request fields such as:
- `width`, `height`
- `num_frames` (video)
- `num_inference_steps`
- `seed`, `fps`
- optional `slo_ms` (per-request SLO target)
By default (when `--dataset-path` is not provided), the script downloads a default trace from
the HuggingFace dataset repo `asukaqaqzz/Dit_Trace`. The default filename can depend on `--task`
(e.g., `t2v` uses a video trace).
Current defaults:
- `--task t2i` -> `sd3_trace.txt`
- `--task t2v` -> `cogvideox_trace.txt`
You can point to your own trace using `--dataset-path`.
## 3. Benchmark Parameters
### Basic flags
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
- `--model`: The OpenAI-compatible `model` field.
- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
- `--num-prompts`: Number of requests to send.
Common optional flags:
- `--output-file`: Write metrics to a JSON file.
- `--disable-tqdm`: Disable the progress bar.
### Resolution / frames / steps: CLI defaults vs dataset fields
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
Precedence rules for `trace` (i.e., what actually gets sent):
- `width/height`: if either `--width` or `--height` is explicitly set, it overrides per-request values from the trace; otherwise per-request values are used when present.
- `num_frames`: per-request `num_frames` takes precedence; otherwise fall back to `--num-frames`.
- `num_inference_steps`: per-request `num_inference_steps` takes precedence; otherwise fall back to `--num-inference-steps`.
### SLO, warmup, and max concurrency
Enable SLO evaluation with `--slo`.
- If a request in the trace already has `slo_ms`, that value is used.
- Otherwise, the script runs warmup requests to infer a base unit time, estimates `expected_ms` by linearly scaling with area/frames/steps, and then sets `slo_ms = expected_ms * --slo-scale`.
Warmup flags:
- `--warmup-requests`: Number of warmup requests.
- `--warmup-num-inference-steps`: Steps used during warmup.
- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
Traffic / concurrency flags:
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
@@ -9,28 +9,29 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Usage:
# Video
t2v:
python3 benchmarks/diffusion_benchmark_serving.py \
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion_benchmark_serving.py \
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task i2v --num-prompts 10
# Image
t2i:
python3 benchmarks/diffusion_benchmark_serving.py \
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
i2i:
python3 benchmarks/diffusion_benchmark_serving.py \
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task i2i --num-prompts 10
"""
import argparse
import ast
import asyncio
import base64
import glob
@@ -40,6 +41,7 @@ import os
import time
import uuid
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field, replace
from typing import Any
@@ -60,6 +62,8 @@ class RequestFuncInput:
num_inference_steps: int | None = None
seed: int | None = None
fps: int | None = None
timestamp: float | None = None
slo_ms: float | None = None
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -73,6 +77,7 @@ class RequestFuncOutput:
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
class BaseDataset(ABC):
@@ -310,6 +315,226 @@ class VBenchDataset(BaseDataset):
return [self[i] for i in range(len(self))]
class TraceDataset(BaseDataset):
"""Trace-based dataset loader for heterogeneous diffusion requests."""
DEFAULT_REPO_ID = "asukaqaqzz/Dit_Trace"
DEFAULT_FILENAME = "sd3_trace.txt"
DEFAULT_FILENAME_BY_TASK: dict[str, str] = {
# Text-to-image traces (e.g., SD3)
"t2i": "sd3_trace.txt",
# Text-to-video traces (e.g., CogVideoX)
"t2v": "cogvideox_trace.txt",
}
def __init__(self, args, api_url: str, model: str):
super().__init__(args, api_url, model)
self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "vllm-omni", "trace")
self.default_filename = self.DEFAULT_FILENAME_BY_TASK.get(getattr(args, "task", ""), self.DEFAULT_FILENAME)
dataset_root = args.dataset_path
if not dataset_root:
dataset_root = self._download_default_trace()
self.items = self._load_items(dataset_root)
@staticmethod
def _coerce_int(x: Any) -> int | None:
if x is None:
return None
if isinstance(x, bool):
return None
if isinstance(x, int):
return x
try:
s = str(x).strip()
if not s:
return None
return int(float(s))
except Exception:
return None
@staticmethod
def _coerce_float(x: Any) -> float | None:
if x is None:
return None
if isinstance(x, float):
return x
if isinstance(x, int):
return float(x)
try:
s = str(x).strip()
if not s:
return None
return float(s)
except Exception:
return None
def _download_default_trace(self) -> str:
"""Download default trace file from HuggingFace Hub if not provided."""
try:
from huggingface_hub import hf_hub_download
except ImportError as exc:
raise ImportError(
"huggingface_hub is required to download the default trace dataset. "
"Install via `pip install huggingface_hub`."
) from exc
os.makedirs(self.cache_dir, exist_ok=True)
return hf_hub_download(
repo_id=self.DEFAULT_REPO_ID,
filename=self.default_filename,
repo_type="dataset",
local_dir=self.cache_dir,
local_dir_use_symlinks=False,
)
def _expand_paths(self, dataset_path: str | None) -> list[str]:
if not dataset_path:
return []
parts = [p.strip() for p in str(dataset_path).split(",") if p.strip()]
paths: list[str] = []
for p in parts:
if any(ch in p for ch in ["*", "?", "["]):
paths.extend(sorted(glob.glob(p)))
elif os.path.isdir(p):
paths.extend(sorted(glob.glob(os.path.join(p, "**", "*.txt"), recursive=True)))
else:
paths.append(p)
seen = set()
unique_paths = []
for p in paths:
if p not in seen:
seen.add(p)
unique_paths.append(p)
return unique_paths
def _parse_trace_file(self, path: str) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
def parse_request_repr_line(line: str) -> dict[str, Any] | None:
text = line.strip()
if not text:
return None
if not (text.startswith("Request(") and text.endswith(")")):
return None
inner = text[len("Request(") : -1]
try:
expr = ast.parse(f"f({inner})", mode="eval")
if not isinstance(expr.body, ast.Call):
return None
call = expr.body
out: dict[str, Any] = {}
for kw in call.keywords:
if kw.arg is None:
continue
out[kw.arg] = ast.literal_eval(kw.value)
return out
except Exception:
return None
# detect first non-empty line to pick parser
first_non_empty = None
with open(path, encoding="utf-8") as f:
for _ in range(50):
pos = f.tell()
line = f.readline()
if not line:
break
if line.strip():
first_non_empty = line.strip()
f.seek(pos)
break
if first_non_empty is None:
return rows
if first_non_empty.startswith("Request("):
with open(path, encoding="utf-8") as f:
for line in f:
parsed = parse_request_repr_line(line)
if isinstance(parsed, dict):
rows.append(parsed)
return rows
# txt fallback: parse Request(...) lines only
with open(path, encoding="utf-8") as f:
for line in f:
parsed = parse_request_repr_line(line)
if isinstance(parsed, dict):
rows.append(parsed)
return rows
def _load_items(self, dataset_root: str) -> list[dict[str, Any]]:
paths = self._expand_paths(dataset_root)
if not paths:
raise ValueError("No trace files found. Provide --dataset-path or rely on default HuggingFace download.")
items: list[dict[str, Any]] = []
for p in paths:
if not os.path.exists(p):
continue
for row in self._parse_trace_file(p):
if isinstance(row, dict):
row = dict(row)
row.setdefault("_source", p)
items.append(row)
if not items:
raise ValueError("Trace dataset is empty after parsing provided paths.")
if self.args.num_prompts is not None:
items = items[: self.args.num_prompts]
return items
def __len__(self) -> int:
return len(self.items)
def __getitem__(self, idx: int) -> RequestFuncInput:
row = self.items[idx]
prompt = row.get("prompt") or row.get("text") or ""
row_height = self._coerce_int(row.get("height"))
row_width = self._coerce_int(row.get("width"))
num_frames = self._coerce_int(row.get("num_frames"))
num_steps = self._coerce_int(row.get("num_inference_steps"))
seed = self._coerce_int(row.get("seed"))
fps = self._coerce_int(row.get("fps"))
timestamp = self._coerce_float(row.get("timestamp"))
slo_ms = self._coerce_float(row.get("slo_ms"))
image_paths = row.get("image_paths")
override_w = self.args.width
override_h = self.args.height
if override_w is not None or override_h is not None:
width = override_w
height = override_h
else:
width = row_width
height = row_height
return RequestFuncInput(
prompt=str(prompt),
api_url=self.api_url,
model=self.model,
width=width,
height=height,
num_frames=num_frames if num_frames is not None else self.args.num_frames,
num_inference_steps=num_steps if num_steps is not None else self.args.num_inference_steps,
seed=seed if seed is not None else self.args.seed,
fps=fps if fps is not None else self.args.fps,
timestamp=timestamp,
slo_ms=slo_ms,
image_paths=image_paths,
request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
)
def get_requests(self) -> list[RequestFuncInput]:
return [self[i] for i in range(len(self))]
class RandomDataset(BaseDataset):
def __init__(self, args, api_url: str, model: str):
self.args = args
@@ -337,6 +562,128 @@ class RandomDataset(BaseDataset):
return [self[i] for i in range(len(self))]
def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
"""Compute expected execution time (ms) based on a base per-step-per-frame unit time.
Assumes linear scaling with pixel area, frame count, and num_inference_steps.
The base unit represents latency for a 16x16 resolution, single frame, single step.
"""
if base_time_ms is None:
return None
width = req.width if req.width is not None else args.width
height = req.height if req.height is not None else args.height
if width is None or height is None:
return None
frames = req.num_frames if req.num_frames is not None else args.num_frames
steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
frame_scale = frames if isinstance(frames, int) and frames > 0 else 1
step_scale = steps if isinstance(steps, int) and steps > 0 else 1
area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
return float(base_time_ms) * area_units * frame_scale * step_scale
def _infer_slo_base_time_ms_from_warmups(
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
args,
) -> float | None:
"""Infer base SLO unit time from warmup requests.
Returns the median base latency (ms) for a 16x16 resolution, single-frame,
single-step request. Only uses warmups that succeeded and have resolvable
width/height.
"""
candidates_ms: list[float] = []
for req, out in warmup_pairs:
if not out.success or out.latency <= 0:
continue
width = req.width if req.width is not None else args.width
height = req.height if req.height is not None else args.height
if width is None or height is None:
continue
frames = req.num_frames if req.num_frames is not None else args.num_frames
steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
frame_scale = int(frames) if isinstance(frames, int) and frames > 0 else 1
step_scale = int(steps) if isinstance(steps, int) and steps > 0 else 1
area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
denom = area_units * float(frame_scale) * float(step_scale)
if denom <= 0:
continue
candidates_ms.append((out.latency * 1000.0) / denom)
if not candidates_ms:
return None
return float(np.median(candidates_ms))
def _populate_slo_ms_from_warmups(
requests_list: list[RequestFuncInput],
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
args,
) -> list[RequestFuncInput]:
"""Populate missing RequestFuncInput.slo_ms using warmup outputs.
- If a request already has slo_ms (e.g., trace-provided), it is kept as-is.
- If any request has slo_ms is None and we can infer base time from warmups,
we estimate each missing request's expected execution time and set:
req.slo_ms = expected_latency_ms * args.slo_scale
Returns updated requests_list.
"""
if not any(req.slo_ms is None for req in requests_list):
return requests_list
base_time_ms = _infer_slo_base_time_ms_from_warmups(warmup_pairs, args)
if base_time_ms is None:
return requests_list
slo_scale = float(getattr(args, "slo_scale", 3.0))
if slo_scale <= 0:
raise ValueError(f"slo_scale must be positive, got {slo_scale}.")
updated: list[RequestFuncInput] = []
for req in requests_list:
if req.slo_ms is not None:
updated.append(req)
continue
expected_ms = _compute_expected_latency_ms_from_base(req, args, base_time_ms)
updated.append(replace(req, slo_ms=(expected_ms * slo_scale) if expected_ms is not None else None))
return updated
async def iter_requests(
requests_list: list[RequestFuncInput],
request_rate: float,
) -> AsyncGenerator[RequestFuncInput, None]:
"""Yield requests using a fixed interval if request_rate is set.
- If request_rate is inf, all requests are yielded immediately (no sleep).
- Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
"""
if request_rate != float("inf"):
if request_rate <= 0:
raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
interval_s = 1.0 / float(request_rate)
for i, req in enumerate(requests_list):
if request_rate != float("inf") and i > 0:
await asyncio.sleep(interval_s)
yield req
def _guess_mime_type(path: str) -> str:
mime, _ = mimetypes.guess_type(path)
return mime or "application/octet-stream"
@@ -415,12 +762,21 @@ async def async_request_chat_completions(
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
def calculate_metrics(
outputs: list[RequestFuncOutput],
total_duration: float,
requests_list: list[RequestFuncInput],
args,
slo_enabled: bool,
):
success_outputs = [o for o in outputs if o.success]
error_outputs = [o for o in outputs if not o.success]
@@ -442,6 +798,29 @@ def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
"peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
}
if slo_enabled:
slo_defined_total = 0
slo_met_success = 0
for req, out in zip(requests_list, outputs):
if req.slo_ms is None:
continue
slo_defined_total += 1
if out.slo_achieved is None:
continue
if out.slo_achieved:
slo_met_success += 1
slo_attain_all = (slo_met_success / slo_defined_total) if slo_defined_total > 0 else 0.0
metrics.update(
{
"slo_attainment_rate": slo_attain_all,
"slo_met_success": slo_met_success,
"slo_scale": getattr(args, "slo_scale", 3.0),
}
)
return metrics
@@ -475,6 +854,8 @@ async def benchmark(args):
if args.dataset == "vbench":
dataset = VBenchDataset(args, api_url, args.model)
elif args.dataset == "trace":
dataset = TraceDataset(args, api_url, args.model)
elif args.dataset == "random":
dataset = RandomDataset(args, api_url, args.model)
else:
@@ -501,6 +882,7 @@ async def benchmark(args):
pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
async with aiohttp.ClientSession() as session:
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
if args.warmup_requests and requests_list:
print(
f"Running {args.warmup_requests} warmup request(s) \
@@ -513,16 +895,20 @@ async def benchmark(args):
warm_req,
num_inference_steps=args.warmup_num_inference_steps,
)
await limited_request_func(warm_req, session, None)
warm_out = await limited_request_func(warm_req, session, None)
warmup_pairs.append((warm_req, warm_out))
if args.slo:
# Prefer trace-provided per-request slo_ms. Only populate when missing.
requests_list = _populate_slo_ms_from_warmups(
requests_list=requests_list,
warmup_pairs=warmup_pairs,
args=args,
)
start_time = time.perf_counter()
tasks = []
for req in requests_list:
if args.request_rate != float("inf"):
# Poisson process: inter-arrival times follow exponential distribution
interval = np.random.exponential(1.0 / args.request_rate)
await asyncio.sleep(interval)
async for req in iter_requests(requests_list=requests_list, request_rate=args.request_rate):
task = asyncio.create_task(limited_request_func(req, session, pbar))
tasks.append(task)
@@ -532,7 +918,7 @@ async def benchmark(args):
pbar.close()
# Calculate metrics
metrics = calculate_metrics(outputs, total_duration)
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
@@ -561,6 +947,12 @@ async def benchmark(args):
print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
if args.slo:
print(f"{'-' * 50}")
print("{:<40} {:<15.2%}".format("SLO Attainment Rate (all):", metrics.get("slo_attainment_rate", 0.0)))
print("{:<40} {:<15}".format("SLO Met (success count):", str(metrics.get("slo_met_success", 0))))
print("{:<40} {:<15}".format("SLO Scale:", str(metrics.get("slo_scale", 3.0))))
if metrics["peak_memory_mb_max"] > 0:
print(f"{'-' * 50}")
print("{:<40} {:<15.2f}".format("Peak Memory Max (MB):", metrics["peak_memory_mb_max"]))
@@ -590,7 +982,7 @@ if __name__ == "__main__":
"--dataset",
type=str,
default="vbench",
choices=["vbench", "random"],
choices=["vbench", "trace", "random"],
help="Dataset to use.",
)
parser.add_argument(
@@ -645,7 +1037,7 @@ if __name__ == "__main__":
parser.add_argument(
"--num-inference-steps",
type=int,
default=None,
default=50,
help="Number of inference steps (for diffusion models).",
)
parser.add_argument(
@@ -656,6 +1048,21 @@ if __name__ == "__main__":
)
parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.")
parser.add_argument(
"--slo",
action="store_true",
help=(
"Enable SLO calculation and reporting. If trace provides per-request slo_ms, it is used. "
"Otherwise, warmup request(s) are used to infer expected execution time assuming linear "
"scaling by resolution, frames, and steps, then slo_ms = expected_time * --slo-scale."
),
)
parser.add_argument(
"--slo-scale",
type=float,
default=3.0,
help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
)
parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
args = parser.parse_args()