From 989161c49ddfc0e39b745fb6d15b0bf3f7cb2310 Mon Sep 17 00:00:00 2001 From: Hongchen Li <65453204+BuckMulligan99@users.noreply.github.com> Date: Wed, 30 Sep 2026 11:38:07 +0800 Subject: [PATCH] [Benchmark] Add MammothModa2 startup/loading breakdown (#7312) Signed-off-by: redearly123 <65453204+redearly123@users.noreply.github.com> Co-authored-by: redearly123 <65453204+redearly123@users.noreply.github.com> --- .gitignore | 1 + benchmarks/README.md | 7 + benchmarks/mammoth_moda2/README.md | 124 +++++++ benchmarks/mammoth_moda2/__init__.py | 2 + benchmarks/mammoth_moda2/bench_startup.py | 338 ++++++++++++++++++ .../mammoth_moda2/bench_storage_scenarios.sh | 67 ++++ benchmarks/mammoth_moda2/parse_startup_log.py | 245 +++++++++++++ benchmarks/mammoth_moda2/raw_load_bench.py | 127 +++++++ .../mammoth_moda2/reference_results.json | 39 ++ recipes/MammothModa2/MammothModa2.md | 9 + tests/benchmarks/test_mammoth_startup.py | 106 ++++++ 11 files changed, 1065 insertions(+) create mode 100644 benchmarks/mammoth_moda2/README.md create mode 100644 benchmarks/mammoth_moda2/__init__.py create mode 100644 benchmarks/mammoth_moda2/bench_startup.py create mode 100755 benchmarks/mammoth_moda2/bench_storage_scenarios.sh create mode 100644 benchmarks/mammoth_moda2/parse_startup_log.py create mode 100644 benchmarks/mammoth_moda2/raw_load_bench.py create mode 100644 benchmarks/mammoth_moda2/reference_results.json create mode 100644 tests/benchmarks/test_mammoth_startup.py diff --git a/.gitignore b/.gitignore index 34253d6cd..ae59008c0 100644 --- a/.gitignore +++ b/.gitignore @@ -243,6 +243,7 @@ datasets/ *.csv *.json !recipes/cosmos3/negative_prompt.json +!benchmarks/mammoth_moda2/reference_results.json !tests/dfx/perf/tests/*.json !apps/ComfyUI-vLLM-Omni/example_workflows/*.json # Omni-DuplexEval CI config JSON (e.g. accuracy/omni_duplex_eval) diff --git a/benchmarks/README.md b/benchmarks/README.md index 1a23b1d48..997d824c1 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -66,6 +66,13 @@ Accuracy benchmarks for image generation/editing models, adapting external suite - **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench) - **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints +### [MammothModa2](mammoth_moda2/README.md) — Startup and Weight Loading + +Startup / model-loading benchmark for the two-stage MammothModa2 (AR → DiT) deployment, plus a raw safetensors loading micro-benchmark. + +- **Layout**: `mammoth_moda2/bench_startup.py` (single-process startup and first/steady request timing), `mammoth_moda2/parse_startup_log.py` (per-stage breakdown from engine logs), `mammoth_moda2/bench_storage_scenarios.sh` (cold/warm page cache on local disk vs. network storage), `mammoth_moda2/raw_load_bench.py` (safetensors → GPU without vLLM) +- **Key metrics**: time to engine ready, per-stage spawn/init/weight-load/profile time, first vs. steady-state request latency + ### Common serving metrics framework `vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include: diff --git a/benchmarks/mammoth_moda2/README.md b/benchmarks/mammoth_moda2/README.md new file mode 100644 index 000000000..3b530d298 --- /dev/null +++ b/benchmarks/mammoth_moda2/README.md @@ -0,0 +1,124 @@ +# MammothModa2 startup benchmark + +Measures `Omni()` initialization and first/subsequent image requests for +MammothModa2-Preview. This is a measurement baseline for +[#7075](https://github.com/vllm-project/vllm-omni/issues/7075); it changes no runtime code. + +## Run + +From the repository root, with the model downloaded and vLLM-Omni installed: + +```bash +set -o pipefail +python benchmarks/mammoth_moda2/bench_startup.py \ + --model /path/MammothModa2-Preview \ + --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ + --height 512 --width 512 --seed 42 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 8}' \ + --repeat 1 --label baseline --output-json baseline.json --save-image baseline.png \ + 2>&1 | tee baseline.log +python benchmarks/mammoth_moda2/parse_startup_log.py baseline.log --require-complete --markdown +``` + +Use `OMP_NUM_THREADS=16` to set CPU threads and `--parallel-stage-init` to +initialize stages concurrently. For first/subsequent requests, use +`--height 1024 --width 1024 --repeat 3` and 50 inference steps. Memory polling +is opt-in (`--sample-memory`); collect it in separate runs to avoid perturbing +the timing comparison. `--save-image` saves only the first request's image. +Compare that image across runs; a fixed seed alone does not guarantee equality +across execution modes or software versions. + +## Measurement boundaries + +| Metric | Boundary / limitation | +| --- | --- | +| Imports | Before torch import → benchmark imports complete; excludes interpreter startup | +| Time to ready | `Omni()` call → return; excludes client imports and requests | +| Time to first image | Benchmark clock start → first `generate()` return, including request setup | +| First / subsequent request | Each `generate()` call → return; same prompt and settings | +| Stage spawn / device init | Log timestamps: launch → `world_size=` → `Starting to load model`; 1 s resolution | +| Weight load | Engine-reported `Loading weights took`; read/cast/copy are merged | +| Model setup estimate | Reported model-load time minus weight-load time; not independently instrumented | +| Profile + KV + warmup | One merged interval after loading; compilation/capture may be nested within it | +| Compile / graph capture | Explicit engine log durations when present; N/A for eager execution | +| Memory | Optional 1 s samples: sum of process-tree RSS and maximum device memory used across GPUs reported by `nvidia-smi`; includes other GPU users and can miss transient peaks | + +Stage timelines overlap with parallel initialization; do not sum them to +calculate time to ready. Missing log fields stay unmeasured. The parser marks +logs without `BENCH_JSON` as incomplete; `--require-complete` makes this an error. + +## Reference results + +A800-SXM4-80GB; Ubuntu 22.04; 16-CPU cgroup quota on a 128-core host; +vLLM 0.28.0+cu129, torch 2.13.0+cu129, vLLM-Omni `34aa1d2a`. +Checkpoint: 34.5 GiB, 8 shards, mostly fp32, on local NVMe. +Committed eager deploy config, GPU memory fractions 0.5 / 0.3. + +The 2026-09-10 rerun used the corrected benchmark: one warm-up process, +then three rounds of four configurations with order rotated between rounds. +The checkpoint page cache was warmed before the matrix; no cache eviction +was requested during it. Each fresh process sent one 512×512, 8-step request +(guidance 4.0, seed 42). Memory polling was disabled. All 12 measured runs +are retained in [the results](reference_results.json). +Values below are `Omni()` seconds, mean ± sample standard deviation, n=3; +client imports took approximately another 10 s. + +| CPU threads | Serial init | Parallel init | +| --- | --- | --- | +| Default (64) | 60.1 ± 0.2 | 41.1 ± 1.0 | +| 16 | 56.5 ± 0.9 | 36.7 ± 0.2 | + +On this configuration, parallel initialization reduced time to ready by about +32%; combining it with 16 threads reduced it by about 39%. This supports +trying these settings on similar deployments, not changing defaults for all +models or hardware. Fixed 512×512 samples matched across all 12 measured runs; +this is a smoke check, not a broad quality evaluation. + +For context, in the earlier serial baseline AR / DiT spent approximately 11 / 11 s in +spawn and imports, 6 / 6 s in device setup, 6 / 4 s loading weights, and +5 / 2 s in profile + KV + warmup. About 6 s followed the last stage becoming +ready. These coarse log intervals locate work; they are not independent +operator timings. + +Separate observations from the earlier experiment (not rerun in this matrix): + +- Requested page-cache eviction: 80.8 s; FlashInfer JIT cache removed: + 129.1 s (means, n=3 each). These are separate cache conditions, not guaranteed + cold-machine startup or isolated disk/JIT costs. +- Three 1024×1024, 50-step eager requests in one process took + 95.2 / 94.9 / 94.6 s. No large first-request penalty was observed in that run. +- With AR `enforce_eager: false` and serial init, time to ready was + 62.5 ± 0.1 s (warm caches, n=3); logged graph capture was 4 s / 0.20 GiB. + The model reported torch.compile unsupported. KV capacity fell from about + 15.7 to 2.0 GiB, and one parallel-init trial failed. The reported non-KV + profiling difference is not a measurement of graph allocation. Graph-mode + samples differed from eager. This needs separate memory and quality analysis + before recommending the configuration. + +## Loading diagnostic and storage scenarios + +```bash +OMP_NUM_THREADS=16 python benchmarks/mammoth_moda2/raw_load_bench.py \ + --model /path/MammothModa2-Preview --shards 6,7,8 --mode cpu_cast +# Compare --mode gpu_cast in a fresh process with the same cache conditions. + +MODEL_LOCAL=/local/MammothModa2-Preview OUT_DIR=./startup-bench \ + bash benchmarks/mammoth_moda2/bench_storage_scenarios.sh +# Optional MODEL_NFS=/nfs/MammothModa2-Preview; SCENARIOS="local-cold local-warm". +``` + +The diagnostic clones each mmap-backed tensor into owned CPU memory before +casting. `materialize_s` includes page faults, CPU allocation and copying; +`cast_s` measures dtype conversion; `h2d_s` includes GPU allocation and copy. +CUDA context/operation warmup is outside the timer. Total also includes file +opening and bookkeeping. This extra CPU copy differs from the engine loader, +so these numbers do not decompose its loading time or measure disk bandwidth. +The revised diagnostic was run three times each for CPU cast with 64 / 16 +threads and GPU cast with 16 threads; per-run values are in the results file. +Both CPU materialization and conversion were faster with 16 threads here; this +does not isolate how much of the engine-level improvement comes from either. + +Storage scenarios explicitly read shards before warm runs and request +`posix_fadvise(DONTNEED)` before cold runs. Eviction is not guaranteed, and +server-side NFS caches are uncontrolled. The script stops on failure and +summarizes only logs from the current invocation. diff --git a/benchmarks/mammoth_moda2/__init__.py b/benchmarks/mammoth_moda2/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/benchmarks/mammoth_moda2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/benchmarks/mammoth_moda2/bench_startup.py b/benchmarks/mammoth_moda2/bench_startup.py new file mode 100644 index 000000000..85093bbdf --- /dev/null +++ b/benchmarks/mammoth_moda2/bench_startup.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Startup / loading benchmark for multi-stage omni models (MammothModa2 first). + +Measures, from a single fresh process: + * python import time (torch + vllm_omni) + * Omni() construction wall time (all stages: weight load, profiling, KV cache, warmup) + * first request latency (includes JIT / lazy init that only happens on first run) + * steady-state request latency (repeat N-1 more requests) + * optional host RSS / device memory samples (--sample-memory, once a second) + +Per-stage breakdown (weight loading, model-load memory, engine init) is emitted by the +stage engine-core subprocesses into the log; use parse_startup_log.py on the captured +stdout/stderr to merge them with the JSON written here. + +Usage (keep the log: the per-stage breakdown is parsed from it afterwards): + python benchmarks/mammoth_moda2/bench_startup.py --model /path/MammothModa2-Preview \ + --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ + --height 1024 --width 1024 --seed 42 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ + --repeat 2 --label nfs-warm --output-json out.json 2>&1 | tee run.log + python benchmarks/mammoth_moda2/parse_startup_log.py run.log --markdown +""" + +import argparse +import json +import os +import platform +import subprocess +import threading +import time + +T_PROCESS_START = time.perf_counter() + +import torch # noqa: E402 + +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args # noqa: E402 +from vllm_omni.entrypoints.omni import Omni # noqa: E402 +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params # noqa: E402 +from vllm_omni.inputs.data import OmniDiffusionSamplingParams # noqa: E402 +from vllm_omni.model_extras import ( # noqa: E402 + build_text_to_image_prompt as build_model_text_to_image_prompt, +) +from vllm_omni.model_extras import ( # noqa: E402 + get_extra_body_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) +from vllm_omni.platforms import current_omni_platform # noqa: E402 + +T_IMPORTS_DONE = time.perf_counter() + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", required=True) + p.add_argument("--deploy-config", default=None) + p.add_argument("--prompt", default="A stylish woman riding a motorcycle in NYC, movie poster style") + p.add_argument("--negative-prompt", default=None) + p.add_argument("--height", type=int, default=1024) + p.add_argument("--width", type=int, default=1024) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--extra-body", type=json.loads, default=None, help="JSON object forwarded as model extra_body") + p.add_argument( + "--repeat", type=int, default=2, help="number of generate() calls; 1st = first-request, rest = steady" + ) + p.add_argument("--label", default="run", help="free-form label, e.g. nfs-cold / nfs-warm / local-nvme") + p.add_argument("--output-json", default=None) + p.add_argument("--save-image", default=None, help="optional path to save the first request's image") + p.add_argument("--sample-memory", action="store_true", help="sample memory once a second (adds overhead)") + p.add_argument( + "--parallel-stage-init", + action="store_true", + help="initialize stages that share a GPU concurrently (VllmOmniOrchestratorConfig.parallel_stage_init)", + ) + args = p.parse_args() + if args.repeat < 1: + p.error("--repeat must be at least 1") + return args + + +def gpu_info() -> dict: + info = {"torch": torch.__version__, "cuda": torch.version.cuda, "python": platform.python_version()} + try: + import vllm + + info["vllm"] = vllm.__version__ + except Exception: # pragma: no cover + pass + try: + import vllm_omni + + info["vllm_omni"] = getattr(vllm_omni, "__version__", "src") + except Exception: # pragma: no cover + pass + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=name,driver_version,memory.total", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=10, + ) + info["gpu"] = out.stdout.strip() + except Exception: + info["gpu"] = "n/a" + return info + + +class PeakSampler: + """Poll host RSS of this process tree and GPU memory in use once a second (benchmark-side only).""" + + def __init__(self, interval_s: float = 1.0): + self.interval_s = interval_s + self.host_rss_peak_gib = 0.0 + self.gpu_used_peak_gib = 0.0 + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _sample(self) -> None: + try: + import psutil + + root = psutil.Process() + rss = root.memory_info().rss + sum( + c.memory_info().rss for c in root.children(recursive=True) if c.is_running() + ) + self.host_rss_peak_gib = max(self.host_rss_peak_gib, rss / 2**30) + except Exception: + pass + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], + capture_output=True, + text=True, + timeout=5, + ) + used = max(float(x) for x in out.stdout.split() if x.strip()) + self.gpu_used_peak_gib = max(self.gpu_used_peak_gib, used / 1024) + except Exception: + pass + + def _run(self) -> None: + while not self._stop.is_set(): + self._sample() + self._stop.wait(self.interval_s) + + def start(self) -> "PeakSampler": + self._thread.start() + return self + + def stop(self) -> dict: + self._stop.set() + self._thread.join(timeout=10) + self._sample() + return { + "host_rss_peak_gib": round(self.host_rss_peak_gib, 2), + "gpu_used_peak_gib": round(self.gpu_used_peak_gib, 2), + } + + +def fs_type(path: str) -> str: + try: + out = subprocess.run(["stat", "-f", "-c", "%T", path], capture_output=True, text=True, timeout=10) + return out.stdout.strip() + except Exception: + return "unknown" + + +def build_request(omni: Omni, args: argparse.Namespace, generator: torch.Generator): + model_class_name = get_model_class_name(omni) + declared = get_extra_body_params(model_class_name) + + prompt_dict = {"prompt": args.prompt} + if args.negative_prompt: + prompt_dict["negative_prompt"] = args.negative_prompt + prompt_dict = build_model_text_to_image_prompt( + model_class_name=model_class_name, prompt=prompt_dict, height=args.height, width=args.width + ) + + diffusion_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + seed=args.seed, + generator=generator, + ) + user_extra = dict(args.extra_body or {}) + if declared: + apply_declared_extra_args(diffusion_params, declared, user_extra) + else: + diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None}) + + init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(model_class_name) + sampling_params_list = [clone_sampling_params(p) for p in (omni.default_sampling_params_list or [])] + if not sampling_params_list: + sampling_params_list = [diffusion_params] + + replaced = False + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + replaced = True + elif init_non_diffusion and hasattr(params, "extra_args"): + params.extra_args = {**(params.extra_args or {}), **(diffusion_params.extra_args or {})} + if hasattr(params, "seed"): + params.seed = args.seed + # MammothModa2 AR stage: one visual token per grid cell + one EOL per row + 1 look-ahead. + info = prompt_dict.get("additional_information", {}) + if idx == 0 and info.get("omni_task") == ["t2i"]: + ar_w = int(info.get("ar_width", [0])[0]) + ar_h = int(info.get("ar_height", [0])[0]) + if ar_w > 0 and ar_h > 0: + params.max_tokens = ar_h * (ar_w + 1) + 1 + if not replaced and len(sampling_params_list) == 1: + sampling_params_list = [diffusion_params] + return prompt_dict, sampling_params_list + + +def stage_metrics(outputs) -> dict: + """Pull vllm-omni's per-stage timing out of the request output when present.""" + res = {} + for out in outputs or []: + metrics = getattr(out, "metrics", None) + if isinstance(metrics, dict) and "stage_metrics" in metrics: + for sid, m in metrics["stage_metrics"].items(): + if not isinstance(m, dict): + continue + res[str(sid)] = { + k: m.get(k) + for k in ("stage_gen_time_ms", "num_tokens_in", "num_tokens_out", "vllm_ttft_ms", "vllm_tpot_ms") + if k in m + } + return res + + +def main() -> None: + args = parse_args() + result = { + "label": args.label, + "model": args.model, + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "sample_memory": args.sample_memory, + "model_fs": fs_type(args.model), + "deploy_config": args.deploy_config, + "height": args.height, + "width": args.width, + "seed": args.seed, + "extra_body": args.extra_body, + "parallel_stage_init": args.parallel_stage_init, + "omp_num_threads": os.environ.get("OMP_NUM_THREADS"), + "env": gpu_info(), + "phases_s": {"imports": round(T_IMPORTS_DONE - T_PROCESS_START, 3)}, + "requests": [], + } + + omni_kwargs = {"model": args.model, "mode": "text-to-image", "log_stats": True} + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + if args.parallel_stage_init: + omni_kwargs["parallel_stage_init"] = True + + sampler = PeakSampler().start() if args.sample_memory else None + t0 = time.perf_counter() + omni = None + try: + omni = Omni(**omni_kwargs) + t1 = time.perf_counter() + result["peak_mem_startup"] = sampler.stop() if sampler else {} + sampler = None + result["phases_s"]["engine_init"] = round(t1 - t0, 3) + result["phases_s"]["process_to_engine_ready"] = round(t1 - T_PROCESS_START, 3) + print( + f"BENCH engine_init={t1 - t0:.2f}s process_to_engine_ready={t1 - T_PROCESS_START:.2f}s " + f"peak_mem_startup={result['peak_mem_startup']}", + flush=True, + ) + sampler = PeakSampler().start() if args.sample_memory else None + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + outputs = None + first_outputs = None + first_image_elapsed = 0.0 + for i in range(args.repeat): + prompt_dict, sampling_params_list = build_request(omni, args, generator) + tr0 = time.perf_counter() + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) + tr1 = time.perf_counter() + if first_outputs is None: + first_outputs = outputs + first_image_elapsed = tr1 - T_PROCESS_START + rec = { + "index": i, + "kind": "first" if i == 0 else "steady", + "latency_s": round(tr1 - tr0, 3), + "stages": stage_metrics(outputs), + } + result["requests"].append(rec) + print(f"BENCH request[{i}] {rec['kind']} latency={tr1 - tr0:.2f}s stages={rec['stages']}", flush=True) + + result["peak_mem_requests"] = sampler.stop() if sampler else {} + sampler = None + steady = [r["latency_s"] for r in result["requests"] if r["kind"] == "steady"] + result["summary_s"] = { + "imports": result["phases_s"]["imports"], + "engine_init": result["phases_s"]["engine_init"], + "first_request": result["requests"][0]["latency_s"], + "steady_request_avg": round(sum(steady) / len(steady), 3) if steady else None, + "time_to_first_image_from_process_start": round(first_image_elapsed, 3), + } + + for key in ("host_rss_peak_gib", "gpu_used_peak_gib"): + result["summary_s"][key] = ( + max(result["peak_mem_startup"][key], result["peak_mem_requests"][key]) if args.sample_memory else None + ) + + from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs + + images = extract_images_from_outputs(first_outputs) + if not images: + raise RuntimeError("First request returned no image") + if args.save_image: + os.makedirs(os.path.dirname(os.path.abspath(args.save_image)), exist_ok=True) + images[0].save(args.save_image) + + print("BENCH_JSON " + json.dumps(result), flush=True) + if args.output_json: + with open(args.output_json, "w") as f: + json.dump(result, f, indent=2) + print(f"BENCH wrote {args.output_json}", flush=True) + finally: + if sampler: + sampler.stop() + if omni is not None: + omni.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/mammoth_moda2/bench_storage_scenarios.sh b/benchmarks/mammoth_moda2/bench_storage_scenarios.sh new file mode 100755 index 000000000..14f56efc9 --- /dev/null +++ b/benchmarks/mammoth_moda2/bench_storage_scenarios.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# MODEL_LOCAL=/local/model MODEL_NFS=/nfs/model OUT_DIR=./startup-bench bash "$0" +# Optional SCENARIOS: local-cold local-warm nfs-cold nfs-warm. +set -euo pipefail +HERE=$(cd "$(dirname "$0")" && pwd) +OUT_DIR=${OUT_DIR:-./startup-bench} +DEPLOY_CONFIG=${DEPLOY_CONFIG:-vllm_omni/deploy/mammoth_moda2.yaml} +STEPS=${STEPS:-50} +SIZE=${SIZE:-1024} +REPEAT=${REPEAT:-2} +SEED=${SEED:-42} +EXTRA=${EXTRA:-'{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": '$STEPS'}'} +mkdir -p "$OUT_DIR/logs" "$OUT_DIR/json" "$OUT_DIR/images" +rm -f "$OUT_DIR/summary.json" "$OUT_DIR/summary.md" + +prepare_cache() { + python - "$1" "$2" <<'PY' +import os +import sys +from pathlib import Path + +files = sorted(Path(sys.argv[1]).glob("*.safetensors")) +if not files: + raise SystemExit("no checkpoint shards found") +for path in files: + with path.open("rb") as f: + if sys.argv[2] == "cold": + os.posix_fadvise(f.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) + else: + while f.read(8 * 1024 * 1024): + pass +print("[bench]", "requested eviction of" if sys.argv[2] == "cold" else "read", len(files), "shards") +PY +} + +default_scenarios="" +if [ -n "${MODEL_LOCAL:-}" ]; then default_scenarios="local-cold local-warm"; fi +if [ -n "${MODEL_NFS:-}" ]; then default_scenarios="$default_scenarios nfs-cold nfs-warm"; fi +read -r -a scenarios <<< "${SCENARIOS:-$default_scenarios}" +if [ "${#scenarios[@]}" -eq 0 ]; then echo "Set MODEL_LOCAL or MODEL_NFS" >&2; exit 1; fi +logs=() +for label in "${scenarios[@]}"; do + case "$label" in + local-cold|local-warm) model=${MODEL_LOCAL:?Set MODEL_LOCAL} ;; + nfs-cold|nfs-warm) model=${MODEL_NFS:?Set MODEL_NFS} ;; + *) echo "Unknown scenario: $label" >&2; exit 1 ;; + esac + prepare_cache "$model" "${label##*-}" + log="$OUT_DIR/logs/$label.log" + # Remove previous outputs for this label so a failed rerun cannot leave stale results. + rm -f "$OUT_DIR/json/$label.json" "$OUT_DIR/images/$label.png" + if python "$HERE/bench_startup.py" --model "$model" --deploy-config "$DEPLOY_CONFIG" \ + --height "$SIZE" --width "$SIZE" --seed "$SEED" --extra-body "$EXTRA" --repeat "$REPEAT" \ + --label "$label" --output-json "$OUT_DIR/json/$label.json" --save-image "$OUT_DIR/images/$label.png" \ + > "$log" 2>&1; then + logs+=("$log") + else + rc=$? + echo "[bench] $label failed (exit $rc); see $log" >&2 + exit "$rc" + fi +done +python "$HERE/parse_startup_log.py" "${logs[@]}" --require-complete --markdown \ + --json-out "$OUT_DIR/summary.json" | tee "$OUT_DIR/summary.md" +echo BENCH_ALL_DONE diff --git a/benchmarks/mammoth_moda2/parse_startup_log.py b/benchmarks/mammoth_moda2/parse_startup_log.py new file mode 100644 index 000000000..b6eb41746 --- /dev/null +++ b/benchmarks/mammoth_moda2/parse_startup_log.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Merge per-stage startup numbers from a bench_startup.py log with its BENCH_JSON line. + +The stage engine cores run in subprocesses and report weight-loading / init timings +only through log lines such as: + + (StageEngineCoreProc_stage0_replica0 pid=1) INFO ... Loading weights took 39.66 seconds + (StageEngineCoreProc_stage0_replica0 pid=1) INFO ... Model loading took 21.45 GiB memory and 40.59 seconds + (StageEngineCoreProc_stage0_replica0 pid=1) INFO ... init engine (profile, create kv cache, warmup model) took 7.00 s + (StageEngineCoreProc_stage0_replica0 pid=1) INFO ... Filesystem type for checkpoints: NFS4. + Checkpoint size: 34.52 GiB. Available RAM: 22.76 GiB. + INFO ... [Omni] AsyncOmniEngine initialized in 159.74 seconds + +Usage: + python benchmarks/mammoth_moda2/parse_startup_log.py run.log [more.log ...] [--markdown] [--json-out summary.json] +""" + +import argparse +import json +import re +import sys +from datetime import datetime +from typing import Any + +STAGE = r"\(StageEngineCoreProc_stage(?P\d+)_replica\d+ pid=\d+\)" +PATTERNS = { + "weights_load_s": re.compile(STAGE + r".*Loading weights took (?P[\d.]+) seconds"), + "model_load_s": re.compile(STAGE + r".*Model loading took (?P[\d.]+) GiB memory and (?P[\d.]+) seconds"), + "engine_init_s": re.compile( + STAGE + r".*init engine \(profile, create kv cache, warmup model\) took (?P[\d.]+) s" + ), + # Only present when the stage does not run with enforce_eager. + "graph_capture_s": re.compile( + STAGE + r".*Graph capturing finished in (?P[\d.]+) secs, took (?P[\d.]+) GiB" + ), + "compile_s": re.compile(STAGE + r".*torch\.compile takes (?P[\d.]+) s in total"), + "ckpt_fs": re.compile( + STAGE + + r".*Filesystem type for checkpoints: (?P\w+)\. Checkpoint size: (?P[\d.]+) GiB\." + + r" Available RAM: (?P[\d.]+) GiB" + ), +} +ENGINE_READY = re.compile(r"AsyncOmniEngine initialized in (?P[\d.]+) seconds") +PREFETCH_SKIP = re.compile(STAGE + r".*exceeds 90% of available RAM") +COMPILE_UNSUPPORTED = re.compile(STAGE + r".*`torch\.compile` is turned on, but the model .* does not support it") + +# Startup timeline: (milestone key, regex). Timestamps in vllm logs are "MM-DD HH:MM:SS" (1 s resolution). +TS = re.compile(r"(?:INFO|WARNING|ERROR) (?P\d\d-\d\d \d\d:\d\d:\d\d)") +MILESTONES = [ + ("omni_init_start", re.compile(r"\[Omni\] Initializing with model")), + ("stage{stage}_launch", re.compile(r"Stage-(?P\d+) set runtime devices")), + ("stage{stage}_proc_up", re.compile(STAGE + r".*world_size=\d+ rank=\d+")), + ("stage{stage}_load_start", re.compile(STAGE + r".*Starting to load model")), + ("stage{stage}_load_done", re.compile(STAGE + r".*Loading weights took")), + ("stage{stage}_init_done", re.compile(STAGE + r".*init engine \(profile, create kv cache, warmup model\) took")), + ("stage{stage}_ready", re.compile(r"\[StageRuntime\] Stage (?P\d+) initialized")), + ("engine_ready", re.compile(r"AsyncOmniEngine initialized in")), +] + + +def _ts_seconds(ts: str) -> int: + # Use a leap year because engine timestamps do not include the year. + value = datetime.strptime("2000-" + ts, "%Y-%m-%d %H:%M:%S") + return int((value - datetime(2000, 1, 1)).total_seconds()) + + +def timeline(lines: list[str]) -> dict: + """Return milestone timestamps (s, relative to omni_init_start) and derived per-stage phase durations.""" + marks: dict[str, int] = {} + for line in lines: + tsm = TS.search(line) + if not tsm: + continue + for key, pat in MILESTONES: + m = pat.search(line) + if m: + k = key.format(stage=m.group("stage")) if "{stage}" in key else key + marks.setdefault(k, _ts_seconds(tsm.group("ts"))) + break + if "omni_init_start" not in marks: + return {} + t0 = marks["omni_init_start"] + rel = {k: (v - t0) % (366 * 86400) for k, v in marks.items()} + phases: dict[str, Any] = {} + names = ("launch", "proc_up", "load_start", "load_done", "init_done", "ready") + for s in sorted({k.split("_")[0] for k in rel if k.startswith("stage")}): + found = {name: rel.get(f"{s}_{name}") for name in names} + if any(v is None for v in found.values()): + continue + ts = {name: int(v) for name, v in found.items() if v is not None} + phases[s] = { + "spawn_import_config_s": ts["proc_up"] - ts["launch"], + "device_dist_init_s": ts["load_start"] - ts["proc_up"], + "model_and_weights_load_s": ts["load_done"] - ts["load_start"], + "profile_kv_warmup_s": ts["init_done"] - ts["load_done"], + "ready_handoff_s": ts["ready"] - ts["init_done"], + "total_s": ts["ready"] - ts["launch"], + } + if "engine_ready" in rel: + last_ready = max((v for k, v in rel.items() if k.endswith("_ready") and k != "engine_ready"), default=0) + phases["post_stages_wiring_s"] = rel["engine_ready"] - last_ready + phases["engine_ready_s"] = rel["engine_ready"] + return {"marks_rel_s": rel, "phases": phases} + + +def parse(path: str) -> dict: + stages: dict[str, dict] = {} + rec: dict[str, Any] = {"log": path, "stages": stages} + with open(path, errors="replace") as f: + lines = f.readlines() + rec["timeline"] = timeline(lines) + for line in lines: + if line.startswith("BENCH_JSON "): + rec["bench"] = json.loads(line[len("BENCH_JSON ") :]) + continue + m = ENGINE_READY.search(line) + if m: + rec["engine_ready_s_reported"] = float(m.group("v")) + continue + m = PREFETCH_SKIP.search(line) + if m: + stages.setdefault(m.group("stage"), {})["page_cache_prefetch_skipped"] = True + continue + m = COMPILE_UNSUPPORTED.search(line) + if m: + stages.setdefault(m.group("stage"), {})["torch_compile_unsupported"] = True + continue + for key, pat in PATTERNS.items(): + m = pat.search(line) + if not m: + continue + st = stages.setdefault(m.group("stage"), {}) + if key == "ckpt_fs": + st["checkpoint_fs"] = m.group("fs") + st["checkpoint_size_gib"] = float(m.group("size")) + st["available_ram_gib"] = float(m.group("ram")) + elif key == "graph_capture_s": + st["graph_capture_s"] = float(m.group("v")) + st["graph_capture_mem_gib"] = float(m.group("mem")) + elif key == "model_load_s": + st["model_load_s"] = float(m.group("v")) + st["model_load_mem_gib"] = float(m.group("mem")) + # "Model loading took" covers model construction + weight loading; "Loading weights took" only + # the weights. The difference also includes loader setup/teardown overhead. + if "weights_load_s" in st: + st["model_setup_estimate_s"] = round(st["model_load_s"] - st["weights_load_s"], 2) + else: + st[key] = float(m.group("v")) + break + rec["status"] = "complete" if "bench" in rec else "incomplete" + return rec + + +def _table(columns: dict[str, str], records: list[dict]) -> str: + def row(values): + return "| " + " | ".join("" if v is None else str(v) for v in values) + " |" + + lines = [row(columns.values()), row(["---"] * len(columns))] + lines.extend(row(record.get(key) for key in columns) for record in records) + return "\n".join(lines) + + +def to_markdown(recs: list[dict]) -> str: + summary, stages, phases, warnings = [], [], [], [] + for r in recs: + b = r.get("bench", {}) + label = b.get("label", r["log"]) + if r.get("status") == "incomplete": + warnings.append(f"Incomplete log: {r['log']} (no BENCH_JSON).") + summary.append({"label": label, **b.get("summary_s", {})}) + for sid, st in sorted(r["stages"].items()): + stages.append({"label": label, "stage": sid, **st}) + if st.get("torch_compile_unsupported"): + stages[-1]["compile_s"] = "unsupported" + ph = r.get("timeline", {}).get("phases", {}) + for sid, values in sorted(ph.items()): + if sid.startswith("stage"): + phases.append({"label": label, "stage": sid, **values}) + if "post_stages_wiring_s" in ph: + phases.append({"label": label, "stage": "after last stage", "total_s": ph["post_stages_wiring_s"]}) + tables = [ + ( + { + "label": "label", + "imports": "imports (s)", + "engine_init": "engine ready (s)", + "first_request": "first request (s)", + "steady_request_avg": "subsequent avg (s)", + "time_to_first_image_from_process_start": "process→first image (s)", + "host_rss_peak_gib": "sampled host RSS (GiB)", + "gpu_used_peak_gib": "sampled GPU used (GiB)", + }, + summary, + ), + ( + { + "label": "label", + "stage": "stage", + "model_setup_estimate_s": "model setup estimate (s)", + "weights_load_s": "weights load (s)", + "compile_s": "compile (s)", + "graph_capture_s": "graph capture (s)", + "engine_init_s": "profile/KV/warmup reported (s)", + }, + stages, + ), + ( + { + "label": "label", + "stage": "stage", + "spawn_import_config_s": "spawn/import (s)", + "device_dist_init_s": "device init (s)", + "model_and_weights_load_s": "model+weights load (s)", + "profile_kv_warmup_s": "profile/KV/warmup interval (s)", + "total_s": "total (s)", + }, + phases, + ), + ] + return "\n\n".join(warnings + [_table(columns, rows) for columns, rows in tables]) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("logs", nargs="+") + ap.add_argument("--markdown", action="store_true") + ap.add_argument("--json-out", default=None) + ap.add_argument("--require-complete", action="store_true", help="fail if any log has no BENCH_JSON") + args = ap.parse_args() + recs = [parse(p) for p in args.logs] + if args.require_complete and any(r["status"] != "complete" for r in recs): + ap.error("incomplete benchmark log: missing BENCH_JSON") + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(recs, f, indent=2) + if args.markdown: + print(to_markdown(recs)) + else: + json.dump(recs, sys.stdout, indent=2) + print() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/mammoth_moda2/raw_load_bench.py b/benchmarks/mammoth_moda2/raw_load_bench.py new file mode 100644 index 000000000..1d67b91a0 --- /dev/null +++ b/benchmarks/mammoth_moda2/raw_load_bench.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Raw safetensors -> GPU micro-benchmark, without vLLM's loader. + +Compares CPU and GPU dtype conversion after materializing each mmap-backed tensor. +Materialization includes page faults, CPU allocation and copying; it is not a disk +bandwidth measurement. CUDA setup is excluded; transfer includes GPU allocation. + + --mode cpu_cast materialize -> cast to bf16 on the CPU -> copy to GPU + --mode gpu_cast materialize -> copy source dtype to GPU -> cast on the GPU + + python benchmarks/mammoth_moda2/raw_load_bench.py --model /path/MammothModa2-Preview --shards 6,7,8 + OMP_NUM_THREADS=4 python benchmarks/mammoth_moda2/raw_load_bench.py --model /path --shards 6,7,8 --mode cpu_cast +""" + +import argparse +import json +import os +import re +import time + +import torch +from safetensors import safe_open + +SHARD_NO = re.compile(r"-(\d+)-of-\d+\.safetensors$") + + +def shard_files(model: str, shards: str) -> list[str]: + """Resolve ``--shards`` to files, refusing empty or partial matches so a typo cannot pass as a fast load.""" + files = sorted(f for f in os.listdir(model) if f.endswith(".safetensors")) + if not files: + raise SystemExit(f"no .safetensors files under {model}") + if shards != "all": + want = {int(x) for x in shards.split(",")} + found = {int(m.group(1)): f for f in files if (m := SHARD_NO.search(f))} + missing = sorted(want - found.keys()) + if missing: + raise SystemExit(f"shard(s) {missing} not found under {model}; available: {sorted(found)}") + files = [found[n] for n in sorted(want)] + return [os.path.join(model, f) for f in files] + + +def sync() -> None: + if torch.cuda.is_available(): + torch.accelerator.synchronize() + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", required=True) + ap.add_argument("--shards", default="all", help="comma list of 1-based shard numbers, or 'all'") + ap.add_argument("--mode", choices=("cpu_cast", "gpu_cast"), default="cpu_cast") + ap.add_argument("--label", default="raw") + ap.add_argument("--no-cuda", action="store_true", help="read + CPU cast only (mode must be cpu_cast)") + a = ap.parse_args() + if a.no_cuda and a.mode != "cpu_cast": + ap.error("--no-cuda only makes sense with --mode cpu_cast") + + paths = shard_files(a.model, a.shards) + total_bytes = sum(os.path.getsize(p) for p in paths) + + if not a.no_cuda: + # Initialize the context and exercise both operations outside the timer. + warmup = torch.zeros(1).to("cuda") + warmup = warmup.to(torch.bfloat16) + sync() + del warmup + + t0 = time.perf_counter() + n_tensors, cpu_bytes = 0, 0 + dtypes: dict[str, int] = {} + t_materialize = t_cast = t_h2d = 0.0 + for p in paths: + with safe_open(p, framework="pt", device="cpu") as f: + for k in f.keys(): + r0 = time.perf_counter() + t = f.get_tensor(k).clone() # touch every page and own the CPU storage + t_materialize += time.perf_counter() - r0 + dtypes[str(t.dtype)] = dtypes.get(str(t.dtype), 0) + 1 + cpu_bytes += t.numel() * t.element_size() + needs_cast = t.is_floating_point() and t.dtype != torch.bfloat16 + if a.mode == "cpu_cast": + c0 = time.perf_counter() + if needs_cast: + t = t.to(torch.bfloat16) + t_cast += time.perf_counter() - c0 + if not a.no_cuda: + h0 = time.perf_counter() + g = t.to("cuda") + sync() + t_h2d += time.perf_counter() - h0 + del g + else: + h0 = time.perf_counter() + g = t.to("cuda") + sync() + t_h2d += time.perf_counter() - h0 + c0 = time.perf_counter() + if needs_cast: + g = g.to(torch.bfloat16) + sync() + t_cast += time.perf_counter() - c0 + del g + n_tensors += 1 + del t + total = time.perf_counter() - t0 + res = { + "label": a.label, + "mode": a.mode, + "files": [os.path.basename(p) for p in paths], + "file_bytes_gib": round(total_bytes / 2**30, 2), + "tensor_bytes_gib": round(cpu_bytes / 2**30, 2), + "n_tensors": n_tensors, + "dtypes": dtypes, + "torch_threads": torch.get_num_threads(), + "cpu_affinity": len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else None, + "materialize_s": round(t_materialize, 3), + "cast_s": round(t_cast, 3), + "h2d_s": round(t_h2d, 3), + "total_s": round(total, 3), + "cuda_setup_included": False, + } + print("RAWLOAD_JSON " + json.dumps(res)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/mammoth_moda2/reference_results.json b/benchmarks/mammoth_moda2/reference_results.json new file mode 100644 index 000000000..de600bfd7 --- /dev/null +++ b/benchmarks/mammoth_moda2/reference_results.json @@ -0,0 +1,39 @@ +{ + "date": "2026-09-10", + "runtime_commit": "34aa1d2afa59547945fc59f88ab6f4f25f52a777", + "environment": {"torch": "2.13.0+cu129", "cuda": "12.9", "python": "3.12.14", "vllm": "0.28.0", "vllm_omni": "0.28.1.dev81+g34aa1d2af.d20260908", "gpu": "NVIDIA A800-SXM4-80GB, 580.126.09, 81920 MiB"}, + "method": "One excluded warm-up/canary; warm page cache; three rounds of the 2x2 eager configurations with rotated order; fresh process for every run; memory polling disabled.", + "request": {"height": 512, "width": 512, "seed": 42, "steps": 8, "guidance": 4.0, "prompt": "A stylish woman riding a motorcycle in NYC, movie poster style"}, + "gpu_memory_utilization": [0.5, 0.3], + "cpu_quota": 16, + "startup_columns": ["configuration", "round", "engine_init_s", "imports_s", "first_request_s", "time_to_first_image_s", "image_md5"], + "startup_runs": [ + ["default-serial", 1, 60.022, 10.158, 20.827, 91.089, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["default-serial", 2, 60.015, 10.082, 20.766, 90.909, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["default-serial", 3, 60.34, 9.911, 21.089, 91.381, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-serial", 1, 55.45, 9.965, 20.711, 86.169, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-serial", 2, 57.01, 9.667, 20.697, 87.418, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-serial", 3, 57.044, 9.707, 20.415, 87.211, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["default-parallel", 1, 40.258, 10.112, 20.392, 70.829, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["default-parallel", 2, 42.177, 10.135, 20.692, 73.093, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["default-parallel", 3, 40.848, 9.997, 20.995, 71.887, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-parallel", 1, 36.88, 9.666, 20.319, 66.932, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-parallel", 2, 36.534, 9.6, 20.314, 66.517, "60c1fe4d464fba8c6ffd1b752fbbd40c"], + ["omp-parallel", 3, 36.689, 9.582, 20.467, 66.802, "60c1fe4d464fba8c6ffd1b752fbbd40c"] + ], + "raw_loading_columns": ["configuration", "round", "torch_threads", "n_tensors", "materialize_s", "cast_s", "h2d_s", "total_s"], + "raw_loading_runs": [ + ["cpu64", 1, 64, 620, 3.174, 2.309, 1.217, 7.796], + ["cpu64", 2, 64, 620, 2.957, 2.443, 1.433, 7.846], + ["cpu64", 3, 64, 620, 2.916, 2.613, 1.067, 7.942], + ["cpu16", 1, 16, 620, 0.857, 1.137, 0.713, 3.643], + ["cpu16", 2, 16, 620, 0.871, 1.174, 0.746, 3.608], + ["cpu16", 3, 16, 620, 0.814, 1.142, 0.724, 3.487], + ["gpu16", 1, 16, 620, 0.938, 0.028, 1.41, 3.74], + ["gpu16", 2, 16, 620, 0.926, 0.029, 1.467, 3.83], + ["gpu16", 3, 16, 620, 0.87, 0.029, 1.395, 3.661] + ], + "raw_loading_workload": {"shards": [6, 7, 8], "source_dtype": "float32", "tensor_gib": 13.3, "warmup": "CUDA context, copy, cast and synchronize before timer", "order_per_round": ["cpu64", "cpu16", "gpu16"]}, + "executed_script_sha256": {"parse_startup_log.py": "aa260de2cfcaec2fc4a678b245bb7c16fa464543aaec72b947fd4f769d547011", "raw_load_bench.py": "1f70f6b3c205402b2c2c7c52b224e1be3c04a12d12d11044a86f1b9b6c7a1d0f", "bench_startup.py": "04b1e8b9398b4d9885fbb6fe341e8a9ee98792d0be200c15c57e1fc5183c6ff0", "__init__.py": "b1d617f45f34dbe6f5f58d792d90c84af8d8d0a7e7def77c48170ca93495ca0d"}, + "notes": ["Standard deviations use ddof=1.", "Only first-request ONGs are compared.", "Raw materialization includes CPU allocation/copy; H2D includes GPU allocation; not engine-loader timings."] +} diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 94084fec5..9b6fab76c 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -289,6 +289,15 @@ The first request took 85.224 seconds. The AR stage generated 4,161 visual token The output was a valid 1024 by 1024 RGB PNG. +### Startup benchmark + +The [startup benchmark](../../benchmarks/mammoth_moda2/README.md) provides +commands, measurement boundaries and a single-A800 eager baseline. On that +16-CPU-quota machine, parallel stage initialization and 16 CPU threads reduced +`Omni()` initialization from about 60 s to 37 s. See the benchmark for the +configuration and limitations; CUDA-graph mode needs separate memory and +quality validation. + ## MammothModa2-Dev unified inference MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses diff --git a/tests/benchmarks/test_mammoth_startup.py b/tests/benchmarks/test_mammoth_startup.py new file mode 100644 index 000000000..2c6f0ab3b --- /dev/null +++ b/tests/benchmarks/test_mammoth_startup.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression checks for benchmark completion and timestamp boundaries (no GPU).""" + +import importlib.util +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] +BENCH = Path(__file__).resolve().parents[2] / "benchmarks" / "mammoth_moda2" +spec = importlib.util.spec_from_file_location("parse_startup_log", BENCH / "parse_startup_log.py") +assert spec and spec.loader +parser = importlib.util.module_from_spec(spec) +spec.loader.exec_module(parser) + + +@pytest.mark.parametrize( + ("start", "end"), + [("09-10 23:59:58", "09-11 00:00:03"), ("12-31 23:59:58", "01-01 00:00:03")], +) +def test_timeline_date_rollover(start, end): + lines = [ + f"INFO {start} [Omni] Initializing with model test", + f"INFO {end} AsyncOmniEngine initialized in 5.00 seconds", + ] + assert parser.timeline(lines)["phases"]["engine_ready_s"] == 5 + + +def test_incomplete_log_is_not_a_success(tmp_path): + log = tmp_path / "failed.log" + log.write_text("RuntimeError: No available memory for the cache blocks\n") + record = parser.parse(str(log)) + assert record["status"] == "incomplete" + assert "Incomplete log" in parser.to_markdown([record]) + result = subprocess.run( + [sys.executable, str(BENCH / "parse_startup_log.py"), str(log), "--require-complete"], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert "missing BENCH_JSON" in result.stderr + + +@pytest.mark.parametrize("exit_code", [0, 23]) +def test_storage_script_completion_and_log_selection(tmp_path, exit_code): + # Copy the wrapper beside a failing engine and the real parser. No GPU or + # cache eviction: local-warm reads one tiny dummy shard before the failure. + for name in ("bench_storage_scenarios.sh", "parse_startup_log.py"): + shutil.copy(BENCH / name, tmp_path / name) + completion = 'BENCH_JSON {"label": "current"}' + (tmp_path / "bench_startup.py").write_text( + f"raise SystemExit({exit_code})\n" if exit_code else f"print({completion!r})\n" + ) + model = tmp_path / "model" + model.mkdir() + (model / "model.safetensors").write_bytes(b"test") + binaries = tmp_path / "bin" + binaries.mkdir() + (binaries / "python").symlink_to(sys.executable) + env = dict( + os.environ, + PATH=str(binaries) + os.pathsep + os.environ["PATH"], + MODEL_LOCAL=str(model), + SCENARIOS="local-warm", + OUT_DIR=str(tmp_path / "out"), + ) + output = tmp_path / "out" + (output / "logs").mkdir(parents=True) + (output / "summary.json").write_text("stale summary") + (output / "logs" / "old.log").write_text('BENCH_JSON {"label": "old"}\n') + result = subprocess.run( + ["bash", str(tmp_path / "bench_storage_scenarios.sh")], env=env, capture_output=True, text=True + ) + assert result.returncode == exit_code + if exit_code: + assert "BENCH_ALL_DONE" not in result.stdout + assert not (output / "summary.json").exists() + else: + summary = json.loads((output / "summary.json").read_text()) + assert [record["bench"]["label"] for record in summary] == ["current"] + + +def test_raw_loader_checks_shards_and_cpu_materialization(tmp_path, monkeypatch, capsys): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + spec = importlib.util.spec_from_file_location("raw_load_bench", BENCH / "raw_load_bench.py") + assert spec and spec.loader + raw = importlib.util.module_from_spec(spec) + spec.loader.exec_module(raw) + safetensors.save_file({"weight": torch.ones(16)}, tmp_path / "model-00006-of-00008.safetensors") + for selection in ("999", "6,999"): + with pytest.raises(SystemExit, match="not found"): + raw.shard_files(str(tmp_path), selection) + assert len(raw.shard_files(str(tmp_path), "6")) == 1 + monkeypatch.setattr(sys, "argv", ["raw", "--model", str(tmp_path), "--shards", "6", "--no-cuda"]) + raw.main() + result = json.loads(capsys.readouterr().out.removeprefix("RAWLOAD_JSON ")) + assert result["n_tensors"] == 1 + assert result["h2d_s"] == 0 + assert "materialize_s" in result and "read_gbps" not in result