From 1f5eca56863d973b4a621b5838174f1fa5e17235 Mon Sep 17 00:00:00 2001 From: Didan Deng <33117903+wtomin@users.noreply.github.com> Date: Tue, 24 Mar 2026 00:19:20 +0800 Subject: [PATCH] [Perf] Qwen-Image Performance Nightly CI test (#1805) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .buildkite/test-nightly.yml | 46 ++ benchmarks/diffusion/backends.py | 257 +++++- .../diffusion/diffusion_benchmark_serving.py | 20 +- tests/perf/scripts/run_diffusion_benchmark.py | 729 ++++++++++++++++++ tests/perf/tests/benchmark_params.json | 57 ++ .../test_qwen_image_sglang_diffusion.json | 73 ++ .../perf/tests/test_qwen_image_vllm_omni.json | 58 ++ vllm_omni/diffusion/data.py | 3 + vllm_omni/diffusion/diffusion_engine.py | 4 + .../worker/diffusion_model_runner.py | 34 + vllm_omni/entrypoints/omni_base.py | 2 + vllm_omni/entrypoints/openai/serving_chat.py | 4 + vllm_omni/outputs.py | 8 + 13 files changed, 1291 insertions(+), 4 deletions(-) create mode 100644 tests/perf/scripts/run_diffusion_benchmark.py create mode 100644 tests/perf/tests/benchmark_params.json create mode 100644 tests/perf/tests/test_qwen_image_sglang_diffusion.json create mode 100644 tests/perf/tests/test_qwen_image_vllm_omni.json diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 2ed96d382..2329800d6 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -246,6 +246,52 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: Qwen-Image Diffusion Perf Test with H100" + key: nightly-qwen-image-performance + timeout_in_minutes: 300 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export DIFFUSION_BENCHMARK_DIR=tests/perf/results + - export CACHE_DIT_VERSION=1.3.0 + - pytest -s -v tests/perf/scripts/run_diffusion_benchmark.py --config-file tests/perf/tests/test_qwen_image_vllm_omni.json + - buildkite-agent artifact upload "tests/perf/results/benchmark_results_*.json" + - buildkite-agent artifact upload "tests/perf/results/logs/*.log" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 4 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":email: Nightly Collection & Email" key: nightly-perf-distribution depends_on: diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index f64b0e350..f12995f6c 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -1,5 +1,6 @@ import asyncio import base64 +import json import mimetypes import os import time @@ -109,8 +110,6 @@ async def async_request_chat_completions( resp_json = await response.json() output.response_body = resp_json output.success = True - if "peak_memory_mb" in resp_json: - output.peak_memory_mb = resp_json["peak_memory_mb"] try: choices = resp_json.get("choices", []) if choices and isinstance(choices, list): @@ -120,7 +119,8 @@ async def async_request_chat_completions( if content and isinstance(content, list) and len(content) > 0: first_item = content[0] if isinstance(first_item, dict): - output.stage_durations = first_item.get("stage_durations") + output.stage_durations = first_item.get("stage_durations") or {} + output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0) except (IndexError, TypeError, AttributeError): pass else: @@ -335,12 +335,263 @@ async def async_request_v1_videos( return output +async def async_request_image_sglang( + input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, +) -> RequestFuncOutput: + output = RequestFuncOutput() + output.start_time = time.perf_counter() + + # Check if we need to use multipart (for image edits with input images) + if input.image_paths and len(input.image_paths) > 0: + # Use multipart/form-data for image edits + data = aiohttp.FormData() + data.add_field("model", input.model) + data.add_field("prompt", input.prompt) + data.add_field("response_format", "b64_json") + + if input.width and input.height: + data.add_field("size", f"{input.width}x{input.height}") + + # Merge extra parameters + for key, value in input.extra_body.items(): + data.add_field(key, str(value)) + + # Add image file(s) + for idx, img_path in enumerate(input.image_paths): + if os.path.exists(img_path): + data.add_field( + "image", + open(img_path, "rb"), + filename=os.path.basename(img_path), + content_type="application/octet-stream", + ) + else: + output.error = f"Image file not found: {img_path}" + output.success = False + if pbar: + pbar.update(1) + return output + + try: + async with session.post(input.api_url, data=data) as response: + if response.status == 200: + resp_json = await response.json() + output.response_body = resp_json + output.success = True + if "peak_memory_mb" in resp_json: + output.peak_memory_mb = resp_json["peak_memory_mb"] + else: + output.error = f"HTTP {response.status}: {await response.text()}" + output.success = False + except Exception as e: + output.error = str(e) + output.success = False + else: + # Use JSON for text-to-image generation + payload = { + "model": input.model, + "prompt": input.prompt, + "n": 1, + "response_format": "b64_json", + } + + 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 + + payload.update(input.extra_body) + + try: + async with session.post(input.api_url, json=payload) as response: + if response.status == 200: + resp_json = await response.json() + output.response_body = resp_json + output.success = True + if "peak_memory_mb" in resp_json: + output.peak_memory_mb = resp_json["peak_memory_mb"] + else: + output.error = f"HTTP {response.status}: {await response.text()}" + output.success = False + except Exception as e: + output.error = str(e) + output.success = False + + output.latency = time.perf_counter() - output.start_time + + # Check SLO if defined + if input.slo_ms is not None and output.success: + output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms + + if pbar: + pbar.update(1) + return output + + +async def async_request_video_sglang( + input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, +) -> RequestFuncOutput: + output = RequestFuncOutput() + output.start_time = time.perf_counter() + + # 1. Submit Job + job_id = None + # Check if we need to upload images (Multipart) or just send JSON + if input.image_paths and len(input.image_paths) > 0: + # Use multipart/form-data + data = aiohttp.FormData() + data.add_field("model", input.model) + data.add_field("prompt", input.prompt) + + if input.width and input.height: + data.add_field("size", f"{input.width}x{input.height}") + + # Add extra body fields to form data if possible, or assume simple key-values + # Note: Nested dicts in extra_body might need JSON serialization if API expects it stringified + if input.extra_body: + data.add_field("extra_body", json.dumps(input.extra_body)) + + # Explicitly add fps/num_frames if they are not in extra_body (bench_serving logic overrides) + if input.num_frames: + data.add_field("num_frames", str(input.num_frames)) + if input.fps: + data.add_field("fps", str(input.fps)) + + # Add image file + # Currently only support single image upload as 'input_reference' per API spec + img_path = input.image_paths[0] + if os.path.exists(img_path): + data.add_field( + "input_reference", + open(img_path, "rb"), + filename=os.path.basename(img_path), + content_type="application/octet-stream", + ) + else: + output.error = f"Image file not found: {img_path}" + output.success = False + if pbar: + pbar.update(1) + return output + + try: + async with session.post(input.api_url, data=data) as response: + if response.status == 200: + resp_json = await response.json() + job_id = resp_json.get("id") + else: + output.error = f"Submit failed HTTP {response.status}: {await response.text()}" + output.success = False + if pbar: + pbar.update(1) + return output + except Exception as e: + output.error = f"Submit exception: {str(e)}" + output.success = False + if pbar: + pbar.update(1) + return output + + else: + # Use JSON + payload: dict[str, Any] = { + "model": input.model, + "prompt": input.prompt, + } + if input.width and input.height: + payload["size"] = f"{input.width}x{input.height}" + if input.num_frames: + payload["num_frames"] = input.num_frames + if input.fps: + payload["fps"] = input.fps + if input.num_inference_steps: + payload["num_inference_steps"] = input.num_inference_steps + + payload.update(input.extra_body) + + try: + async with session.post(input.api_url, json=payload) as response: + if response.status == 200: + resp_json = await response.json() + job_id = resp_json.get("id") + else: + output.error = f"Submit failed HTTP {response.status}: {await response.text()}" + output.success = False + if pbar: + pbar.update(1) + return output + except Exception as e: + output.error = f"Submit exception: {str(e)}" + output.success = False + if pbar: + pbar.update(1) + return output + + if not job_id: + output.error = "No job_id returned" + output.success = False + if pbar: + pbar.update(1) + return output + + # 2. Poll for completion + # Assuming the API returns a 'status' field. + # We construct the check URL. Assuming api_url is like .../v1/videos + # The check url should be .../v1/videos/{id} + check_url = f"{input.api_url}/{job_id}" + + while True: + try: + async with session.get(check_url) as response: + if response.status == 200: + status_data = await response.json() + status = status_data.get("status") + if status == "completed": + output.success = True + output.response_body = status_data + if "peak_memory_mb" in status_data: + output.peak_memory_mb = status_data["peak_memory_mb"] + break + elif status == "failed": + output.success = False + output.error = f"Job failed: {status_data.get('error')}" + break + else: + # queued or processing + await asyncio.sleep(1.0) + else: + output.success = False + output.error = f"Poll failed HTTP {response.status}: {await response.text()}" + break + except Exception as e: + output.success = False + output.error = f"Poll exception: {str(e)}" + break + + output.latency = time.perf_counter() - output.start_time + + # Check SLO if defined + if input.slo_ms is not None and output.success: + output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms + + if pbar: + pbar.update(1) + return output + + backends_function_mapping = { "2i": { "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"), "openai": (async_request_openai_images, "/v1/images/generations"), + "sglang": (async_request_image_sglang, "/v1/images/generations"), }, "2v": { "v1/videos": (async_request_v1_videos, "/v1/videos"), + "sglang": (async_request_video_sglang, "/v1/videos"), }, } diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 602f9938e..7178742d3 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -745,6 +745,15 @@ def calculate_metrics( latencies = [o.latency for o in success_outputs] peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0] + # Aggregate per-stage durations across all successful requests that reported them. + stage_duration_lists: dict[str, list[float]] = {} + for o in success_outputs: + for stage, duration in (o.stage_durations or {}).items(): + stage_duration_lists.setdefault(stage, []).append(duration) + stage_durations_mean = {s: float(np.mean(v)) for s, v in stage_duration_lists.items()} + stage_durations_p50 = {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()} + stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()} + metrics = { "duration": total_duration, "completed_requests": num_success, @@ -758,6 +767,9 @@ def calculate_metrics( "peak_memory_mb_max": max(peak_memories) if peak_memories else 0, "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0, "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0, + "stage_durations_mean": stage_durations_mean, + "stage_durations_p50": stage_durations_p50, + "stage_durations_p99": stage_durations_p99, } if slo_enabled: @@ -953,6 +965,12 @@ async def benchmark(args): print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"])) print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"])) + if metrics["stage_durations_mean"]: + print(f"{'-' * 50}") + print("Stage Durations Mean (s):") + for stage, val in metrics["stage_durations_mean"].items(): + print("{:<40} {:<15.4f}".format(f" {stage}:", val)) + print("\n" + "=" * 60) if args.output_file: @@ -976,7 +994,7 @@ if __name__ == "__main__": "--backend", type=str, default="vllm-omni", - choices=["vllm-omni", "openai", "v1/videos"], + choices=["vllm-omni", "openai", "sglang", "v1/videos"], help="Backend to target the benchmark to.", ) parser.add_argument( diff --git a/tests/perf/scripts/run_diffusion_benchmark.py b/tests/perf/scripts/run_diffusion_benchmark.py new file mode 100644 index 000000000..3fb49f6c3 --- /dev/null +++ b/tests/perf/scripts/run_diffusion_benchmark.py @@ -0,0 +1,729 @@ +""" +Performance benchmark CI runner for diffusion models. + +Supports two server backends: + - vllm-omni (default): starts DiffusionServer via vllm_omni.entrypoints.cli.main, + benchmarks with diffusion_benchmark_serving.py --backend vllm-omni + - sglang: starts SglangServer via `sglang serve`, + benchmarks with diffusion_benchmark_serving.py --backend sglang + +A config JSON file is REQUIRED via --config-file: + pytest run_diffusion_benchmark.py --config-file tests/perf/tests/test_qwen_image_vllm_omni.json + pytest run_diffusion_benchmark.py --config-file tests/perf/tests/test_qwen_image_sglang_diffusion.json + +JSON config entries are distinguished by a "server_type" field ("vllm-omni" or "sglang"). +sglang entries support two additional fields under server_params: + - "env": dict of extra environment variables (e.g. SGLANG_CACHE_DIT_ENABLED) + - "cache_dit_config": dict written to a temp YAML and passed as + --cache-dit-config to sglang serve (requires cache-dit == 1.3.0) + +All benchmark results for a session are consolidated into a single JSON file under +BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable). +Each entry in the file contains the test metadata (test_name, backend, benchmark_params, +timestamp) together with the raw metrics returned by the benchmark script. +""" + +import importlib.metadata +import json +import os +import socket +import subprocess +import sys +import tempfile +import threading +import time +from datetime import datetime +from pathlib import Path +from typing import Any + +import psutil +import pytest + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + +_DEFAULT_RESULT_DIR = Path(__file__).parent.parent / "results" +BENCHMARK_RESULT_DIR = Path(os.environ.get("DIFFUSION_BENCHMARK_DIR", str(_DEFAULT_RESULT_DIR))) + +BENCHMARK_SCRIPT = str( + Path(__file__).parent.parent.parent.parent / "benchmarks" / "diffusion" / "diffusion_benchmark_serving.py" +) + +# Single aggregated result file for the entire benchmark session. +# Populated lazily after CONFIG_FILE_PATH is resolved. +_SESSION_TIMESTAMP = datetime.now().strftime("%Y%m%d-%H%M%S") +_RESULT_LOCK = threading.Lock() + + +def _get_config_file_from_argv() -> str | None: + """Read --config-file from sys.argv at import time so pytest parametrize can use it. + + pytest_addoption (below) registers the same flag so pytest does not reject it. + Supports both ``--config-file path`` and ``--config-file=path`` forms. + Returns None if the flag is not present; callers must handle the missing case. + """ + for i, arg in enumerate(sys.argv): + if arg == "--config-file" and i + 1 < len(sys.argv): + return sys.argv[i + 1] + if arg.startswith("--config-file="): + return arg.split("=", 1)[1] + return None + + +CONFIG_FILE_PATH = _get_config_file_from_argv() +if CONFIG_FILE_PATH is None: + raise ValueError( + "--config-file is required. Pass the path to a benchmark config JSON, e.g.:\n" + " pytest run_diffusion_benchmark.py " + "--config-file tests/perf/tests/test_qwen_image_vllm_omni.json" + ) + +# --------------------------------------------------------------------------- +# Config loading +# --------------------------------------------------------------------------- + + +def _resolve_refs(configs: list[dict[str, Any]], config_dir: Path) -> list[dict[str, Any]]: + """Resolve {"$ref": "filename.json"} in benchmark_params fields.""" + for cfg in configs: + bp = cfg.get("benchmark_params") + if isinstance(bp, dict) and "$ref" in bp: + ref_path = config_dir / bp["$ref"] + try: + with open(ref_path, encoding="utf-8") as f: + cfg["benchmark_params"] = json.load(f) + except FileNotFoundError: + raise ValueError(f"benchmark_params $ref not found: {ref_path}") + except json.JSONDecodeError as e: + raise ValueError(f"JSON parsing error in {ref_path}: {e}") + return configs + + +def load_configs(config_path: str) -> list[dict[str, Any]]: + try: + abs_path = Path(config_path).resolve() + with open(abs_path, encoding="utf-8") as f: + configs = json.load(f) + return _resolve_refs(configs, abs_path.parent) + except json.JSONDecodeError as e: + raise ValueError(f"JSON parsing error: {str(e)}") + except FileNotFoundError: + raise ValueError(f"Configuration file not found: {config_path}") + except Exception as e: + raise RuntimeError(f"Failed to load configuration file: {str(e)}") + + +BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) + +_config_stem = Path(CONFIG_FILE_PATH).stem # e.g. "test_qwen_image_vllm_omni" +AGGREGATED_RESULT_FILE = BENCHMARK_RESULT_DIR / f"benchmark_results_{_config_stem}_{_SESSION_TIMESTAMP}.json" + + +def _append_to_aggregated_file(record: dict[str, Any]) -> None: + """Thread-safe append of *record* to the session-level aggregated JSON file. + + The file contains a JSON array; each call loads the existing array (or + starts a new one), appends the record, and writes the file back atomically. + """ + with _RESULT_LOCK: + BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True) + if AGGREGATED_RESULT_FILE.exists(): + with open(AGGREGATED_RESULT_FILE, encoding="utf-8") as f: + records: list[dict] = json.load(f) + else: + records = [] + records.append(record) + with open(AGGREGATED_RESULT_FILE, "w", encoding="utf-8") as f: + json.dump(records, f, indent=2, ensure_ascii=False) + + +# Register --config-file with pytest so it does not reject the argument. +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--config-file", + action="store", + default=None, + help=( + "Path to the benchmark config JSON file (required). " + "Example: --config-file tests/perf/tests/test_qwen_image_vllm_omni.json" + ), + ) + + +_server_lock = threading.Lock() + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _get_open_port() -> int: + """Return an available TCP port on localhost.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return s.getsockname()[1] + + +def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None: + """Block until the given host:port accepts connections or timeout expires.""" + start = time.time() + while time.time() - start < timeout: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(1) + if s.connect_ex((host, port)) == 0: + return + except Exception: + pass + time.sleep(2) + raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s") + + +def _kill_process_tree(pid: int) -> None: + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + all_pids = [pid] + [c.pid for c in children] + + for child in children: + try: + child.terminate() + except psutil.NoSuchProcess: + pass + + gone, alive = psutil.wait_procs(children, timeout=10) + for child in alive: + try: + child.kill() + except psutil.NoSuchProcess: + pass + + try: + parent.terminate() + parent.wait(timeout=10) + except (psutil.NoSuchProcess, psutil.TimeoutExpired): + try: + parent.kill() + except psutil.NoSuchProcess: + pass + + time.sleep(1) + still_alive = [p for p in all_pids if psutil.pid_exists(p)] + if still_alive: + print(f"Warning: processes still alive after shutdown: {still_alive}") + for p in still_alive: + try: + subprocess.run(["kill", "-9", str(p)], timeout=2) + except Exception: + pass + except psutil.NoSuchProcess: + pass + + +# --------------------------------------------------------------------------- +# Server classes +# --------------------------------------------------------------------------- + + +class DiffusionServer: + """Start a vLLM-Omni diffusion model server as a subprocess. + + Launched via vllm_omni.entrypoints.cli.main with the diffusion-specific + parallelism flags (--usp, --ring, --cfg-parallel-size, etc.) passed directly + on the CLI. Minimum hardware: 4× NVIDIA H100 80 GB. + """ + + server_type = "vllm-omni" + + def __init__( + self, + model: str, + serve_args: list[str], + *, + port: int | None = None, + ) -> None: + self.model = model + self.serve_args = serve_args + self.host = "127.0.0.1" + self.port = port if port is not None else _get_open_port() + self.proc: subprocess.Popen | None = None + self.test_name: str = "" + + def _start_server(self) -> None: + env = os.environ.copy() + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--omni", + "--host", + self.host, + "--port", + str(self.port), + ] + self.serve_args + + print(f"Launching DiffusionServer: {' '.join(cmd)}") + self.proc = subprocess.Popen( + cmd, + env=env, + cwd=str(Path(__file__).parent.parent.parent.parent), + ) + _wait_for_port(self.host, self.port) + print(f"DiffusionServer ready on {self.host}:{self.port}") + + def __enter__(self): + self._start_server() + return self + + def __exit__(self, *_): + if self.proc: + _kill_process_tree(self.proc.pid) + + +_CACHE_DIT_REQUIRED_VERSION = os.environ.get("CACHE_DIT_VERSION", "1.3.0") + + +def _check_cache_dit_version(required: str = _CACHE_DIT_REQUIRED_VERSION) -> None: + """Verify that the installed cache-dit package matches *required* exactly. + + Raises RuntimeError if the package is not installed or the version differs. + """ + try: + installed = importlib.metadata.version("cache-dit") + except importlib.metadata.PackageNotFoundError: + raise RuntimeError( + f"cache-dit is not installed. Please install version {required}: pip install cache-dit=={required}" + ) + if installed != required: + raise RuntimeError( + f"cache-dit version mismatch: required {required}, " + f"but found {installed}. " + f"Please install the correct version: pip install cache-dit=={required}" + ) + + +class SglangServer: + """Start a sglang serve process for diffusion benchmarking. + + Supports two Cache-DiT activation modes: + 1. Environment variable: pass env={"SGLANG_CACHE_DIT_ENABLED": "true"} + 2. YAML config file: pass cache_dit_config={...} (written to a temp + file and forwarded as --cache-dit-config; requires cache-dit >= 1.3.0) + """ + + server_type = "sglang" + + def __init__( + self, + model: str, + serve_args: list[str], + *, + port: int | None = None, + env_overrides: dict[str, str] | None = None, + cache_dit_config: dict[str, Any] | None = None, + ) -> None: + self.model = model + self.serve_args = serve_args + self.host = "127.0.0.1" + self.port = port if port is not None else _get_open_port() + self.env_overrides = env_overrides or {} + self.cache_dit_config = cache_dit_config + self.proc: subprocess.Popen | None = None + self._tmp_yaml: str | None = None + self.test_name: str = "" + if self.cache_dit_config is not None: + _check_cache_dit_version() + + @staticmethod + def _write_cache_dit_yaml(config: dict[str, Any]) -> str: + """Serialize config dict to a temp YAML file and return its path. + + Tries PyYAML first for clean block-style output; falls back to + json.dump since JSON is valid YAML and correctly handles arbitrary + nesting, lists, booleans, and null values. + """ + tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) + try: + import yaml # PyYAML + + yaml.dump(config, tmp, default_flow_style=False, allow_unicode=True) + except ImportError: + json.dump(config, tmp, indent=2, ensure_ascii=False) + tmp.write("\n") + tmp.close() + print(f" Cache-DiT config written to: {tmp.name}") + return tmp.name + + def _start_server(self) -> None: + env = os.environ.copy() + env.update(self.env_overrides) + + cmd = [ + "sglang", + "serve", + "--model-path", + self.model, + "--host", + self.host, + "--port", + str(self.port), + ] + self.serve_args + + if self.cache_dit_config is not None: + self._tmp_yaml = self._write_cache_dit_yaml(self.cache_dit_config) + cmd += ["--cache-dit-config", self._tmp_yaml] + + print(f"Launching SglangServer: {' '.join(cmd)}") + if self.env_overrides: + print(f" Extra env: {self.env_overrides}") + + self.proc = subprocess.Popen( + cmd, + env=env, + cwd=str(Path(__file__).parent.parent.parent.parent), + ) + _wait_for_port(self.host, self.port) + print(f"SglangServer ready on {self.host}:{self.port}") + + def __enter__(self): + self._start_server() + return self + + def __exit__(self, *_): + if self.proc: + _kill_process_tree(self.proc.pid) + if self._tmp_yaml: + try: + Path(self._tmp_yaml).unlink(missing_ok=True) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Config helpers +# --------------------------------------------------------------------------- + + +def _build_serve_args(serve_args_dict: dict[str, Any], server_type: str = "vllm-omni") -> list[str]: + """Convert a serve_args dict from test.json into a flat CLI argument list. + + Boolean handling differs by server type: + - vllm-omni uses store_true/store_false style: True → add flag only, + False → omit flag entirely. + - sglang accepts explicit boolean values: always emit ``--flag true/false``. + """ + args: list[str] = [] + for key, value in serve_args_dict.items(): + flag = f"--{key}" + if isinstance(value, bool): + if server_type == "sglang": + args.extend([flag, str(value).lower()]) + elif value: + args.append(flag) + elif isinstance(value, dict): + args.extend([flag, json.dumps(value, separators=(",", ":"))]) + else: + args.extend([flag, str(value)]) + return args + + +def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Return one server-config dict per unique test_name.""" + seen: set[str] = set() + result: list[dict[str, Any]] = [] + for cfg in configs: + test_name = cfg["test_name"] + if test_name in seen: + continue + seen.add(test_name) + server_type = cfg.get("server_type", "vllm-omni") + result.append( + { + "test_name": test_name, + "server_type": server_type, + "model": cfg["server_params"]["model"], + "serve_args": _build_serve_args(cfg["server_params"].get("serve_args", {}), server_type), + "env_overrides": cfg["server_params"].get("env", {}), + "cache_dit_config": cfg["server_params"].get("cache_dit_config"), + "benchmark_backend": server_type, # "vllm-omni" or "sglang" + } + ) + return result + + +def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]: + mapping: dict[str, list[dict]] = {} + for cfg in configs: + name = cfg["test_name"] + mapping.setdefault(name, []) + mapping[name].extend(cfg["benchmark_params"]) + return mapping + + +def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer | SglangServer: + """Factory: return the appropriate server instance for the given config.""" + model = server_cfg["model"] + serve_args = server_cfg["serve_args"] + if server_cfg["server_type"] == "sglang": + return SglangServer( + model=model, + serve_args=serve_args, + env_overrides=server_cfg.get("env_overrides", {}), + cache_dit_config=server_cfg.get("cache_dit_config"), + ) + return DiffusionServer(model=model, serve_args=serve_args) + + +# --------------------------------------------------------------------------- +# Parametrize data +# --------------------------------------------------------------------------- + +server_params = _unique_server_params(BENCHMARK_CONFIGS) +test_param_map = _test_param_mapping(BENCHMARK_CONFIGS) + +benchmark_indices: list[int] = list(range(max(len(v) for v in test_param_map.values()))) + +# --------------------------------------------------------------------------- +# Pytest fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def diffusion_server(request): + """Start one server (vllm-omni or sglang) per unique test configuration.""" + with _server_lock: + server_cfg: dict[str, Any] = request.param + test_name = server_cfg["test_name"] + server_type = server_cfg["server_type"] + + print(f"\nStarting {server_type} server for test: {test_name}") + with _make_server(server_cfg) as server: + server.test_name = test_name + print(f"{server_type} server started successfully") + yield server + print(f"{server_type} server stopping…") + + print(f"{server_type} server stopped") + + +@pytest.fixture +def benchmark_params(request, diffusion_server): + """Yield the benchmark params dict for the current (server, index) pair.""" + param_index: int = request.param + test_name = diffusion_server.test_name + + params_list = test_param_map.get(test_name, []) + if not params_list: + raise ValueError(f"No benchmark params for test: {test_name}") + if param_index >= len(params_list): + pytest.skip(f"Param index {param_index} out of range for {test_name} (has {len(params_list)} params)") + + current = param_index + 1 + total = len(params_list) + print(f"\n Running benchmark {current}/{total} for {test_name}") + return {"test_name": test_name, "params": params_list[param_index]} + + +# --------------------------------------------------------------------------- +# Benchmark runner +# --------------------------------------------------------------------------- + + +def run_benchmark( + host: str, + port: int, + model: str, + params: dict[str, Any], + test_name: str, + backend: str = "vllm-omni", +) -> dict[str, Any]: + """Run diffusion_benchmark_serving.py as a subprocess and return parsed metrics. + + The raw metrics are written to a temporary file by the subprocess. After + the run completes the metrics are merged with full metadata (test_name, + backend, benchmark_params, timestamp) and appended to the session-wide + aggregated JSON file (AGGREGATED_RESULT_FILE). The temporary file is + removed afterwards. Subprocess stdout/stderr are tee'd to a .log file + under BENCHMARK_RESULT_DIR/logs/; its path is stored in the record. + """ + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + + log_dir = BENCHMARK_RESULT_DIR / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + log_file = log_dir / f"{test_name}_{backend}_{timestamp}.log" + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp: + tmp_result_file = Path(tmp.name) + + exclude_keys = {"baseline", "dataset", "task", "name"} + + cmd = [ + sys.executable, + BENCHMARK_SCRIPT, + "--host", + host, + "--port", + str(port), + "--model", + model, + "--backend", + backend, + "--dataset", + params.get("dataset", "random"), + "--task", + params.get("task", "t2i"), + "--output-file", + str(tmp_result_file), + ] + + for key, value in params.items(): + if key in exclude_keys or value is None: + continue + flag = f"--{key}" + if isinstance(value, bool): + if value: + cmd.append(flag) + elif isinstance(value, (dict, list)): + cmd.extend([flag, json.dumps(value, separators=(",", ":"))]) + else: + cmd.extend([flag, str(value)]) + + # Insert -u so the subprocess runs with unbuffered stdout/stderr, ensuring + # all print() output is flushed to the pipe immediately instead of being + # held in Python's internal block-buffer until process exit (which can + # cause truncated or out-of-order log output when stdout is piped). + cmd = [cmd[0], "-u"] + cmd[1:] + + print(f"\nRunning benchmark (backend={backend}): {' '.join(cmd)}") + print(f" Log file: {log_file}") + + # Redirect stdout + stderr directly to the log file at the OS level + # (equivalent to `cmd > log 2>&1`), so no output is ever lost regardless + # of how the subprocess exits. The log is echoed to the terminal afterwards. + with open(log_file, "w", encoding="utf-8") as log_fh: + log_fh.write(f"cmd: {' '.join(cmd)}\n\n") + log_fh.flush() + + process = subprocess.Popen( + cmd, + stdout=log_fh, + stderr=log_fh, + cwd=str(Path(__file__).parent.parent.parent.parent), + ) + process.wait() + + with open(log_file, encoding="utf-8") as log_fh: + print(log_fh.read(), end="") + + if process.returncode != 0: + tmp_result_file.unlink(missing_ok=True) + raise RuntimeError(f"Benchmark script exited with code {process.returncode}") + + if not tmp_result_file.exists(): + raise FileNotFoundError(f"Benchmark result file not found: {tmp_result_file}") + + try: + with open(tmp_result_file, encoding="utf-8") as f: + metrics: dict[str, Any] = json.load(f) + finally: + tmp_result_file.unlink(missing_ok=True) + + record: dict[str, Any] = { + "test_name": test_name, + "backend": backend, + "timestamp": timestamp, + "benchmark_params": params, + "result": metrics, + "log_file": str(log_file), + } + _append_to_aggregated_file(record) + print(f"\n Result appended to: {AGGREGATED_RESULT_FILE}") + print(f" Log saved to: {log_file}") + + return metrics + + +# --------------------------------------------------------------------------- +# Assertions +# --------------------------------------------------------------------------- + + +def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None: + """Assert that benchmark metrics satisfy the configured baselines.""" + num_prompts = params.get("num-prompts", 10) + completed = result.get("completed_requests", result.get("completed", 0)) + assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}" + + for metric, threshold in params.get("baseline", {}).items(): + current = result.get(metric) + assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}" + if "throughput" in metric: + assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}" + else: + assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}" + + +# --------------------------------------------------------------------------- +# Test entry point +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "diffusion_server", + server_params, + ids=[p["test_name"] for p in server_params], + indirect=True, +) +@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) +def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): + """Run the diffusion performance benchmark and assert against baselines. + + One server is started per unique parallel configuration (module scope). + For each server, all benchmark parameter sets defined in the config JSON + are executed sequentially; results are asserted against the baselines. + + Tracked metrics: + - throughput_qps (higher is better) + - latency_p50, latency_p99 (lower is better) + """ + test_name = benchmark_params["test_name"] + params = benchmark_params["params"] + backend = diffusion_server.server_type # "vllm-omni" or "sglang" + + result = run_benchmark( + host=diffusion_server.host, + port=diffusion_server.port, + model=diffusion_server.model, + params=params, + test_name=test_name, + backend=backend, + ) + + print(f"\n{'=' * 60}") + print(f"Results for {test_name} (server={diffusion_server.server_type}, backend={backend}):") + for key in ( + "throughput_qps", + "latency_mean", + "latency_median", + "latency_p50", + "latency_p99", + "peak_memory_mb_max", + "peak_memory_mb_mean", + "peak_memory_mb_median", + ): + if key in result: + print(f" {key}: {result[key]:.4f}") + + print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}") + print("=" * 60) + + assert_result(result, params) diff --git a/tests/perf/tests/benchmark_params.json b/tests/perf/tests/benchmark_params.json new file mode 100644 index 000000000..561e778e9 --- /dev/null +++ b/tests/perf/tests/benchmark_params.json @@ -0,0 +1,57 @@ +[ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.001, + "latency_p99": 1000.0, + "peak_memory_mb_max": 400000, + "peak_memory_mb_mean": 400000 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.001, + "latency_p99": 1000.0, + "peak_memory_mb_max": 400000, + "peak_memory_mb_mean": 400000 + } + }, + { + "name": "mixed_resolution", + "dataset": "random", + "task": "t2i", + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "random-request-config": [ + {"width": 512, "height": 512, "num_inference_steps": 20, "weight": 0.15}, + {"width": 768, "height": 768, "num_inference_steps": 20, "weight": 0.25}, + {"width": 1024, "height": 1024, "num_inference_steps": 25, "weight": 0.45}, + {"width": 1536, "height": 1536, "num_inference_steps": 35, "weight": 0.15} + ], + "baseline": { + "throughput_qps": 0.001, + "latency_p99": 1000.0, + "peak_memory_mb_max": 400000, + "peak_memory_mb_mean": 400000 + } + } +] diff --git a/tests/perf/tests/test_qwen_image_sglang_diffusion.json b/tests/perf/tests/test_qwen_image_sglang_diffusion.json new file mode 100644 index 000000000..3d3aa41c6 --- /dev/null +++ b/tests/perf/tests/test_qwen_image_sglang_diffusion.json @@ -0,0 +1,73 @@ +[ + { + "test_name": "test_sglang_qwen_image_single_device", + "description": "sglang: single-device baseline (no parallelism)", + "server_type": "sglang", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "num-gpus": 1, + "dit-cpu-offload": false, + "text-encoder-cpu-offload": false, + "vae-cpu-offload": false + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + }, + + { + "test_name": "test_sglang_qwen_image_ulysses2_cfg2_vae_patch4", + "description": "sglang: USP=2 + CFG-parallel=2 + VAE patch parallel via --use-parallel-tiling", + "server_type": "sglang", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "num-gpus": 4, + "sp-degree": 2, + "ulysses-degree": 2, + "enable-cfg-parallel": true, + "use-parallel-tiling": true, + "dit-cpu-offload": false, + "text-encoder-cpu-offload": false, + "vae-cpu-offload": false + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + }, + + { + "test_name": "test_sglang_qwen_image_ulysses2_cfg2_cache_dit", + "description": "sglang: USP=2 + CFG-parallel=2 + Cache-DiT via YAML config (--cache-dit-config)", + "server_type": "sglang", + "server_params": { + "model": "Qwen/Qwen-Image", + "cache_dit_config": { + "cache_config": { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": false, + "taylorseer_order": 1, + "scm_steps_mask_policy": "none", + "scm_steps_policy": "dynamic" + }, + "parallelism_config": { + "ulysses_size": 2, + "attention_backend": "native" + } + }, + "serve_args": { + "num-gpus": 4, + "sp-degree": 2, + "ulysses-degree": 2, + "enable-cfg-parallel": true, + "dit-cpu-offload": false, + "text-encoder-cpu-offload": false, + "vae-cpu-offload": false + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + } +] diff --git a/tests/perf/tests/test_qwen_image_vllm_omni.json b/tests/perf/tests/test_qwen_image_vllm_omni.json new file mode 100644 index 000000000..12ad4bb3c --- /dev/null +++ b/tests/perf/tests/test_qwen_image_vllm_omni.json @@ -0,0 +1,58 @@ +[ + { + "test_name": "test_qwen_image_single_device", + "description": "Single-device baseline (no parallelism)", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + }, + + { + "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4", + "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + }, + + { + "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit", + "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "cache-backend": "cache_dit", + "cache-config": { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": false, + "taylorseer_order": 1, + "scm_steps_mask_policy": null, + "scm_steps_policy": "dynamic" + }, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": {"$ref": "benchmark_params.json"} + } +] diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index f15aae6a0..70b3d8d13 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -651,6 +651,9 @@ class DiffusionOutput: # logged duration of stages stage_durations: dict[str, float] = field(default_factory=dict) + # memory usage info + peak_memory_mb: float = 0.0 + class AttentionBackendEnum(enum.Enum): FA = enum.auto() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index d51a6cec3..1b36bb0ef 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -180,6 +180,7 @@ class DiffusionEngine: multimodal_output={"audio": request_audio_payload}, final_output_type="audio", stage_durations=output.stage_durations, + peak_memory_mb=output.peak_memory_mb, ), ] else: @@ -196,6 +197,7 @@ class DiffusionEngine: custom_output=output.custom_output or {}, multimodal_output=mm_output, stage_durations=output.stage_durations, + peak_memory_mb=output.peak_memory_mb, ), ] else: @@ -226,6 +228,7 @@ class DiffusionEngine: multimodal_output={"audio": request_audio_payload}, final_output_type="audio", stage_durations=output.stage_durations, + peak_memory_mb=output.peak_memory_mb, ), ) else: @@ -252,6 +255,7 @@ class DiffusionEngine: custom_output=output.custom_output or {}, multimodal_output=mm_output, stage_durations=output.stage_durations, + peak_memory_mb=output.peak_memory_mb, ), ) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 356700f92..32ea5bf64 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -190,6 +190,33 @@ class DiffusionModelRunner: """Load weights into the pipeline.""" return self.pipeline.load_weights(weights) + def _record_peak_memory(self, output: DiffusionOutput) -> None: + """Record peak GPU memory for the current forward pass into output. + + Must be called immediately after pipeline.forward(), with + reset_peak_memory_stats() called just before it, so the measurement + reflects this request only and not the global historical maximum. + + Uses max_memory_reserved (CUDA memory pool high-water mark) rather than + max_memory_allocated so that allocator fragmentation is also visible. + See: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html + """ + peak_reserved_bytes = current_omni_platform.max_memory_reserved() + peak_allocated_bytes = current_omni_platform.max_memory_allocated() + + output.peak_memory_mb = peak_reserved_bytes / (1024**2) + peak_reserved_gb = peak_reserved_bytes / (1024**3) + peak_allocated_gb = peak_allocated_bytes / (1024**3) + pool_overhead_gb = peak_reserved_gb - peak_allocated_gb + + logger.info( + "Peak GPU memory (this request): %.2f GB reserved, %.2f GB allocated, %.2f GB pool overhead (%.1f%%)", + peak_reserved_gb, + peak_allocated_gb, + pool_overhead_gb, + pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0, + ) + def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput: """ Execute a forward pass for the given requests. @@ -239,10 +266,17 @@ class DiffusionModelRunner: ): self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps) + is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + if is_primary: + current_omni_platform.reset_peak_memory_stats() + with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): with record_function("pipeline_forward"): output = self.pipeline.forward(req) + if is_primary: + self._record_peak_memory(output) + # NOTE: if ( self.cache_backend is not None diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ba77f83f4..f64165ae5 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -246,6 +246,7 @@ class OmniBase: req_id = result.get("request_id") engine_outputs = result.get("engine_outputs") stage_durations = getattr(result["engine_outputs"], "stage_durations", {}) + peak_memory_mb = getattr(result["engine_outputs"], "peak_memory_mb", 0.0) finished = engine_outputs.finished submit_ts = result.get("stage_submit_ts") @@ -281,6 +282,7 @@ class OmniBase: request_output=engine_outputs, images=images, stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, ) def shutdown(self) -> None: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 49fd0c03a..39828e95c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -1916,6 +1916,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Handle profiling data stage_durations = omni_outputs.stage_durations + peak_memory_mb = omni_outputs.peak_memory_mb # Handle different image output formats images = [] @@ -1971,6 +1972,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "url": f"data:image/png;base64,{img_base64}", }, "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, } ) @@ -2187,6 +2189,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Handle nested OmniRequestOutput structure where images might be in request_output images = getattr(result.request_output, "images", []) stage_durations = result.stage_durations + peak_memory_mb = result.peak_memory_mb # Convert images to base64 content image_contents: list[dict[str, Any]] = [] @@ -2208,6 +2211,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "url": f"data:image/png;base64,{img_base64}", }, "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, } ) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 10b1e50ed..ca3ba271a 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -65,6 +65,9 @@ class OmniRequestOutput: # profiling data stage_durations: dict[str, float] = field(default_factory=dict) + # memory usage info + peak_memory_mb: float = 0.0 + @classmethod def from_pipeline( cls, @@ -102,6 +105,7 @@ class OmniRequestOutput: custom_output: dict[str, Any] | None = None, final_output_type: str = "image", stage_durations: dict[str, float] | None = None, + peak_memory_mb: float = 0.0, ) -> "OmniRequestOutput": """Create output from diffusion model. @@ -113,6 +117,8 @@ class OmniRequestOutput: latents: Optional latent tensors multimodal_output: Optional multimodal output dict custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds) + stage_durations: Optional stage durations (execution time of each stage) dict + peak_memory_mb: Peak memory usage in MB Returns: OmniRequestOutput configured for diffusion mode @@ -127,6 +133,7 @@ class OmniRequestOutput: _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {}, stage_durations=stage_durations or {}, + peak_memory_mb=peak_memory_mb, finished=True, ) @@ -277,6 +284,7 @@ class OmniRequestOutput: f"multimodal_output={self._multimodal_output}", f"custom_output={self._custom_output}", f"stage_durations={self.stage_durations}", + f"peak_memory_mb={self.peak_memory_mb}", ] return f"OmniRequestOutput({', '.join(parts)})"