[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

@@ -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)

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:

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.

benchmarks/mammoth_moda2/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

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()

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

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<stage>\d+)_replica\d+ pid=\d+\)"
+PATTERNS = {
+    "weights_load_s": re.compile(STAGE + r".*Loading weights took (?P<v>[\d.]+) seconds"),
+    "model_load_s": re.compile(STAGE + r".*Model loading took (?P<mem>[\d.]+) GiB memory and (?P<v>[\d.]+) seconds"),
+    "engine_init_s": re.compile(
+        STAGE + r".*init engine \(profile, create kv cache, warmup model\) took (?P<v>[\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<v>[\d.]+) secs, took (?P<mem>[\d.]+) GiB"
+    ),
+    "compile_s": re.compile(STAGE + r".*torch\.compile takes (?P<v>[\d.]+) s in total"),
+    "ckpt_fs": re.compile(
+        STAGE
+        + r".*Filesystem type for checkpoints: (?P<fs>\w+)\. Checkpoint size: (?P<size>[\d.]+) GiB\."
+        + r" Available RAM: (?P<ram>[\d.]+) GiB"
+    ),
+}
+ENGINE_READY = re.compile(r"AsyncOmniEngine initialized in (?P<v>[\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<ts>\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<stage>\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<stage>\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()

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()

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."]
+}

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

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