[Edge][WP2] Add initialization timeline and cold-start profiler

* vllm_omni/engine/init_timeline.py: InitTimeline phase collector; engine
  phases (stages_init, g3_device_lock, g2_spawn/ready) and worker phases
  (proc_entry, engine_core_init, init_device, load_model/load_weights,
  talker_mtp_init, attn_prewarm, cudagraph_capture, talker_mtp_capture,
  code2wav_cudagraph, predictor_setup_compile/warmup/cudagraph) recorded via
  worker_phase/timed_phase; workers append JSONL under
  VLLM_OMNI_INIT_TIMELINE_DIR (inherited at spawn) and the engine merges them.
  Enabled by VLLM_OMNI_INIT_TIMELINE=<json> or log_stats; table logged at INFO,
  partial table on orchestrator startup timeout, suggested timeouts logged.
* omni_base.py: named deploy profiles default to stage_init_timeout=900 /
  init_timeout=1500 unless given explicitly.
* gpu_model_runner._init_talker_mtp tolerates max_cudagraph_capture_size=None;
  gpu_ar_model_runner.shutdown skips CUDA graph teardown on CPU.
* benchmarks/tts/init_profile.py: runs stream_latency_bench --init-only per
  configuration (default, cold_cache, eager_stage1, eager_both,
  parallel_stage_init, edge, ...) and tabulates phase durations.
* stream_latency_bench: --parallel-stage-init / --deploy-profile flags; the
  resource poller no longer shadows Thread._stop and skips GPU sampling on CPU.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

benchmarks/tts/init_profile.py

@@ -0,0 +1,287 @@
+#!/usr/bin/env python3
+"""Initialization (cold-start) profiler for multi-stage TTS pipelines.
+
+Runs ``stream_latency_bench.py --init-only`` once per configuration as a
+subprocess with ``VLLM_OMNI_INIT_TIMELINE`` set, then tabulates the phase
+durations reported by the init timeline (engine + worker processes).
+
+Configurations are config-only variants of one deploy YAML (flattened through
+``resolve_deploy_yaml`` so ``base_config`` inheritance is preserved) plus
+environment changes for cache warm/cold experiments. Run sequentially on ONE
+GPU through the scheduler, e.g.:
+
+  gpu run --gpus 1 --timeout 40m --note "init_profile warm set" -- env HF_HUB_OFFLINE=1 \\
+    python benchmarks/tts/init_profile.py --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\
+    --configs default,eager_stage1,eager_both,parallel_stage_init,edge
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import datetime as _dt
+import json
+import os
+import subprocess
+import sys
+import tempfile
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+BENCH_DIR = Path(__file__).resolve().parent
+REPO_ROOT = BENCH_DIR.parent.parent
+DEPLOY_DIR = REPO_ROOT / "vllm_omni" / "deploy"
+PHASE_COLUMNS = [
+    "stages_init",
+    "g2_spawn_and_ready",
+    "engine_core_init",
+    "init_device",
+    "load_model",
+    "load_weights",
+    "predictor_setup_compile",
+    "predictor_warmup",
+    "cudagraph_capture",
+    "talker_mtp_capture",
+    "code2wav_cudagraph",
+]
+
+
+@dataclass
+class Config:
+    name: str
+    deploy_updates: dict[str, Any] = field(default_factory=dict)  # {"stages": {idx: {...}}, "top": {...}}
+    env: dict[str, str] = field(default_factory=dict)
+    extra_args: list[str] = field(default_factory=list)
+    deploy_config: str | None = None  # explicit YAML instead of the model default
+    fresh_caches: bool = False
+
+
+CONFIGS: dict[str, Config] = {
+    "default": Config("default"),
+    "cold_cache": Config("cold_cache", fresh_caches=True),
+    "eager_stage1": Config("eager_stage1", deploy_updates={"stages": {1: {"enforce_eager": True}}}),
+    "eager_both": Config(
+        "eager_both", deploy_updates={"stages": {0: {"enforce_eager": True}, 1: {"enforce_eager": True}}}
+    ),
+    "parallel_stage_init": Config("parallel_stage_init", extra_args=["--parallel-stage-init"]),
+    "edge": Config("edge", deploy_config=str(DEPLOY_DIR / "edge" / "qwen3_tts.yaml")),
+    "edge_eager_stage1": Config(
+        "edge_eager_stage1",
+        deploy_config=str(DEPLOY_DIR / "edge" / "qwen3_tts.yaml"),
+        deploy_updates={"stages": {1: {"enforce_eager": True}}},
+    ),
+    "edge_parallel": Config(
+        "edge_parallel", deploy_config=str(DEPLOY_DIR / "edge" / "qwen3_tts.yaml"), extra_args=["--parallel-stage-init"]
+    ),
+}
+
+
+def _flatten_deploy_yaml(path: Path) -> dict[str, Any]:
+    """Resolve ``base_config`` inheritance without importing torch-heavy modules."""
+    with open(path, encoding="utf-8") as f:
+        raw = yaml.safe_load(f) or {}
+    base = raw.pop("base_config", None)
+    if not base:
+        return raw
+    base_path = Path(base) if os.path.isabs(base) else (path.parent / base)
+    merged = _flatten_deploy_yaml(base_path.resolve())
+    for key, value in raw.items():
+        if key == "stages":
+            by_id = {int(s.get("stage_id", i)): dict(s) for i, s in enumerate(merged.get("stages", []))}
+            for s in value or []:
+                sid = int(s.get("stage_id", 0))
+                by_id.setdefault(sid, {"stage_id": sid}).update(s)
+            merged["stages"] = [by_id[k] for k in sorted(by_id)]
+        elif key == "platforms" and isinstance(value, dict) and isinstance(merged.get("platforms"), dict):
+            for plat, block in value.items():
+                merged["platforms"].setdefault(plat, {})
+                merged["platforms"][plat].update(block or {})
+        else:
+            merged[key] = value
+    return merged
+
+
+def write_variant_yaml(base_path: Path, updates: dict[str, Any], out_dir: Path, name: str) -> Path:
+    cfg = copy.deepcopy(_flatten_deploy_yaml(base_path))
+    for key, value in (updates.get("top") or {}).items():
+        cfg[key] = value
+    for idx, stage_updates in (updates.get("stages") or {}).items():
+        for stage in cfg.get("stages", []):
+            if int(stage.get("stage_id", -1)) == int(idx):
+                stage.update(stage_updates)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out = out_dir / f"{name}.yaml"
+    with open(out, "w", encoding="utf-8") as f:
+        yaml.safe_dump(cfg, f, sort_keys=False)
+    return out
+
+
+def phase_durations(timeline: dict[str, Any] | None) -> dict[str, float]:
+    """Sum durations per phase name (a phase may appear per stage/worker)."""
+    out: dict[str, float] = {}
+    if not timeline:
+        return out
+    for p in timeline.get("phases", []):
+        out[p["name"]] = out.get(p["name"], 0.0) + float(p.get("dur_s", 0.0))
+    return out
+
+
+def phase_durations_by_stage(timeline: dict[str, Any] | None) -> dict[str, dict[str, float]]:
+    out: dict[str, dict[str, float]] = {}
+    if not timeline:
+        return out
+    for p in timeline.get("phases", []):
+        key = "-" if p.get("stage_id") is None else str(p["stage_id"])
+        out.setdefault(key, {})
+        out[key][p["name"]] = out[key].get(p["name"], 0.0) + float(p.get("dur_s", 0.0))
+    return out
+
+
+def render_table(rows: list[dict[str, Any]]) -> str:
+    cols = ["config", "init_s", "ttfa_ms", "gpu_mib", *PHASE_COLUMNS]
+    header = "| " + " | ".join(cols) + " |"
+    sep = "|" + "|".join(["---"] * len(cols)) + "|"
+    lines = [header, sep]
+    for r in rows:
+        vals = [
+            r["config"],
+            f"{r.get('init_s', float('nan')):.1f}",
+            f"{(r.get('ttfa_ms') or float('nan')):.1f}",
+            str(r.get("gpu_mib", "")),
+        ]
+        for c in PHASE_COLUMNS:
+            v = r.get("phases", {}).get(c)
+            vals.append("" if v is None else f"{v:.1f}")
+        lines.append("| " + " | ".join(vals) + " |")
+    return "\n".join(lines)
+
+
+def run_config(
+    cfg: Config, model: str, out_dir: Path, timeout_s: int, bench_passthrough: list[str], python: str
+) -> dict[str, Any]:
+    run_dir = out_dir / cfg.name
+    run_dir.mkdir(parents=True, exist_ok=True)
+    env = dict(os.environ)
+    env.update(cfg.env)
+    env["VLLM_OMNI_INIT_TIMELINE"] = str(run_dir / "timeline.json")
+    if cfg.fresh_caches:
+        for var in ("VLLM_CACHE_ROOT", "TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"):
+            env[var] = tempfile.mkdtemp(prefix=f"omni_{var.lower()}_")
+    deploy = cfg.deploy_config or _default_deploy_for(model)
+    if cfg.deploy_updates:
+        deploy = str(write_variant_yaml(Path(deploy), cfg.deploy_updates, run_dir, cfg.name))
+    cmd = [
+        python,
+        str(BENCH_DIR / "stream_latency_bench.py"),
+        "--init-only",
+        "--warmup",
+        "1",
+        "--model",
+        model,
+        "--deploy-config",
+        deploy,
+        "--output-dir",
+        str(run_dir / "bench"),
+        "--tag",
+        cfg.name,
+        *cfg.extra_args,
+        *bench_passthrough,
+    ]
+    t0 = time.perf_counter()
+    with open(run_dir / "stdout.log", "w") as log:
+        proc = subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT, timeout=timeout_s, cwd=str(REPO_ROOT))
+    wall = time.perf_counter() - t0
+    result: dict[str, Any] = {
+        "config": cfg.name,
+        "returncode": proc.returncode,
+        "wall_s": wall,
+        "cmd": cmd,
+        "deploy_config": deploy,
+    }
+    timeline = None
+    tl_path = run_dir / "timeline.json"
+    if tl_path.exists():
+        timeline = json.loads(tl_path.read_text())
+    bench_files = sorted((run_dir / "bench").rglob("*.json"))
+    if bench_files:
+        doc = json.loads(bench_files[-1].read_text())
+        result["init_s"] = doc["init"]["init_s"]
+        reqs = doc.get("requests") or []
+        result["ttfa_ms"] = reqs[0].get("ttfa_ms") if reqs else None
+        result["gpu_mib"] = doc.get("resources", {}).get("peak_gpu_mem_mib")
+        timeline = timeline or doc["init"].get("init_timeline")
+    result["phases"] = phase_durations(timeline)
+    result["phases_by_stage"] = phase_durations_by_stage(timeline)
+    result["total_s"] = timeline.get("total_s") if timeline else None
+    return result
+
+
+def _default_deploy_for(model: str) -> str:
+    # Avoid importing vllm_omni here (heavy); mirror the deploy naming rule for Qwen3-TTS.
+    if "qwen3-tts" in model.lower() or "qwen3_tts" in model.lower():
+        return str(DEPLOY_DIR / "qwen3_tts.yaml")
+    raise SystemExit(f"Pass --deploy-config-base for model {model!r} (no default known here)")
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice")
+    p.add_argument("--configs", default="default,eager_stage1,eager_both,parallel_stage_init,edge")
+    p.add_argument("--repeat", type=int, default=1)
+    p.add_argument("--timeout-s", type=int, default=1500)
+    p.add_argument("--output-dir", default=str(BENCH_DIR / "edge_results" / "init_profile"))
+    p.add_argument("--python", default=sys.executable)
+    p.add_argument("--list", action="store_true")
+    args, passthrough = p.parse_known_args(argv)
+    if args.list:
+        for name, cfg in CONFIGS.items():
+            print(
+                f"{name}: updates={cfg.deploy_updates} env={cfg.env} args={cfg.extra_args} "
+                f"deploy={cfg.deploy_config} fresh={cfg.fresh_caches}"
+            )
+        return 0
+    names = [n.strip() for n in args.configs.split(",") if n.strip()]
+    unknown = [n for n in names if n not in CONFIGS]
+    if unknown:
+        raise SystemExit(f"Unknown configs {unknown}; known: {sorted(CONFIGS)}")
+    ts = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
+    out_dir = Path(args.output_dir) / ts
+    out_dir.mkdir(parents=True, exist_ok=True)
+    rows: list[dict[str, Any]] = []
+    # Default passthrough: generous timeouts so cold compiles finish; callers may override.
+    if "--init-timeout" not in passthrough:
+        passthrough += ["--init-timeout", "1500"]
+    if "--stage-init-timeout" not in passthrough:
+        passthrough += ["--stage-init-timeout", "1200"]
+    for name in names:
+        for rep in range(args.repeat):
+            cfg = CONFIGS[name]
+            if args.repeat > 1:
+                cfg = Config(
+                    f"{name}_r{rep}", cfg.deploy_updates, cfg.env, cfg.extra_args, cfg.deploy_config, cfg.fresh_caches
+                )
+            print(f"[init_profile] running {cfg.name} ...", flush=True)
+            try:
+                row = run_config(cfg, args.model, out_dir, args.timeout_s, passthrough, args.python)
+            except subprocess.TimeoutExpired:
+                row = {"config": cfg.name, "returncode": -1, "error": f"timeout after {args.timeout_s}s", "phases": {}}
+            rows.append(row)
+            print(
+                f"[init_profile] {cfg.name}: rc={row.get('returncode')} init_s={row.get('init_s')} "
+                f"phases={row.get('phases')}",
+                flush=True,
+            )
+    summary = {"model": args.model, "timestamp": ts, "rows": rows, "table": render_table(rows)}
+    with open(out_dir / "summary.json", "w") as f:
+        json.dump(summary, f, indent=1)
+    print(render_table(rows))
+    print(f"[init_profile] wrote {out_dir / 'summary.json'}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/tts/stream_latency_bench.py

@@ -133,6 +133,8 @@ class ResourcePoller(threading.Thread):
         return total
 
     def _sample_gpu(self, pids: set[int]) -> int | None:
+        if os.environ.get("VLLM_TARGET_DEVICE", "").lower() == "cpu":
+            return None  # CPU platform: no GPU belongs to this run
         try:
             out = subprocess.run(
                 ["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],