From 073d31b9cd2cde8f50f69d4b7bdd60fcddfe6ded Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 11:01:39 +0800 Subject: [PATCH] [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= 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 --- benchmarks/tts/init_profile.py | 287 +++++++++++++++ benchmarks/tts/stream_latency_bench.py | 2 + tests/benchmarks/test_init_profile.py | 60 +++ tests/engine/test_init_timeline.py | 110 ++++++ vllm_omni/engine/init_timeline.py | 347 ++++++++++++++++++ vllm_omni/engine/stage_engine_core_proc.py | 18 +- vllm_omni/engine/stage_runtime.py | 35 +- vllm_omni/entrypoints/omni_base.py | 13 +- .../models/common/qwen3_code_predictor.py | 4 + .../models/qwen3_tts/qwen3_tts_code2wav.py | 2 + vllm_omni/worker/gpu_ar_model_runner.py | 17 +- vllm_omni/worker/gpu_ar_worker.py | 2 + vllm_omni/worker/gpu_generation_worker.py | 2 + vllm_omni/worker/gpu_model_runner.py | 10 +- 14 files changed, 885 insertions(+), 24 deletions(-) create mode 100644 benchmarks/tts/init_profile.py create mode 100644 tests/benchmarks/test_init_profile.py create mode 100644 tests/engine/test_init_timeline.py create mode 100644 vllm_omni/engine/init_timeline.py diff --git a/benchmarks/tts/init_profile.py b/benchmarks/tts/init_profile.py new file mode 100644 index 000000000..a60d63811 --- /dev/null +++ b/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()) diff --git a/benchmarks/tts/stream_latency_bench.py b/benchmarks/tts/stream_latency_bench.py index 1a9b9306d..0b5780a19 100644 --- a/benchmarks/tts/stream_latency_bench.py +++ b/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"], diff --git a/tests/benchmarks/test_init_profile.py b/tests/benchmarks/test_init_profile.py new file mode 100644 index 000000000..8e1c8b2b6 --- /dev/null +++ b/tests/benchmarks/test_init_profile.py @@ -0,0 +1,60 @@ +"""Tests for benchmarks/tts/init_profile.py helpers (CPU).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import yaml + +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "benchmarks" / "tts")) +import init_profile as ip # noqa: E402 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_flatten_edge_yaml_inherits_base_and_applies_updates(tmp_path): + edge = ip.DEPLOY_DIR / "edge" / "qwen3_tts.yaml" + flat = ip._flatten_deploy_yaml(edge) + assert "base_config" not in flat + assert [s["stage_id"] for s in flat["stages"]] == [0, 1] + # inherited from base + assert flat["stages"][0]["trust_remote_code"] is True + # overridden by edge + assert flat["stages"][0]["max_num_seqs"] == 4 + assert flat["connectors"]["connector_of_shared_memory"]["extra"]["codec_chunk_ramp"] == [2, 4, 8, 16, 25] + + out = ip.write_variant_yaml( + edge, {"stages": {1: {"enforce_eager": True}}, "top": {"async_chunk": False}}, tmp_path, "v" + ) + doc = yaml.safe_load(out.read_text()) + assert doc["stages"][1]["enforce_eager"] is True + assert doc["stages"][0].get("enforce_eager") is None + assert doc["async_chunk"] is False + + +def test_phase_durations_and_table(): + timeline = { + "total_s": 100.0, + "phases": [ + {"name": "load_weights", "dur_s": 10.0, "stage_id": 0}, + {"name": "load_weights", "dur_s": 5.0, "stage_id": 1}, + {"name": "predictor_setup_compile", "dur_s": 40.0, "stage_id": 0}, + ], + } + d = ip.phase_durations(timeline) + assert d == {"load_weights": 15.0, "predictor_setup_compile": 40.0} + by_stage = ip.phase_durations_by_stage(timeline) + assert by_stage["0"]["load_weights"] == 10.0 and by_stage["1"]["load_weights"] == 5.0 + table = ip.render_table([{"config": "default", "init_s": 100.0, "ttfa_ms": 30.0, "gpu_mib": 5000, "phases": d}]) + assert "| default | 100.0 | 30.0 | 5000 |" in table + assert "predictor_setup_compile" in table.splitlines()[0] + assert ip.phase_durations(None) == {} + + +def test_config_registry_names(): + for name in ("default", "cold_cache", "eager_stage1", "eager_both", "parallel_stage_init", "edge"): + assert name in ip.CONFIGS + assert ip.CONFIGS["cold_cache"].fresh_caches is True + assert ip.CONFIGS["parallel_stage_init"].extra_args == ["--parallel-stage-init"] diff --git a/tests/engine/test_init_timeline.py b/tests/engine/test_init_timeline.py new file mode 100644 index 000000000..7d37a1edb --- /dev/null +++ b/tests/engine/test_init_timeline.py @@ -0,0 +1,110 @@ +"""Tests for vllm_omni/engine/init_timeline.py (CPU, no engine).""" + +from __future__ import annotations + +import json +import os +import time + +import pytest + +from vllm_omni.engine import init_timeline as it + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _reset_timeline(monkeypatch): + it.InitTimeline.reset() + monkeypatch.delenv(it.ENV_DIR, raising=False) + monkeypatch.delenv(it.ENV_STAGE_ID, raising=False) + monkeypatch.delenv(it.ENV_REPLICA_ID, raising=False) + yield + it.InitTimeline.reset() + + +def test_disabled_timeline_records_nothing(): + tl = it.InitTimeline.start(enabled=False) + with tl.phase("x"): + pass + tl.mark("m") + assert tl.phases == [] + assert it.ENV_DIR not in os.environ + with it.worker_phase("w"): + pass # no-op without env / enabled timeline + + +def test_phase_nesting_and_durations(tmp_path): + tl = it.InitTimeline.start(enabled=True, run_dir=tmp_path / "run") + with tl.phase("outer", stage_id=0) as info: + time.sleep(0.01) + with tl.phase("inner", stage_id=0, n=3): + time.sleep(0.01) + info["k"] = "v" + names = [p.name for p in tl.phases] + assert names == ["inner", "outer"] # inner closes first + outer = tl.phases[1] + assert outer.dur_s >= 0.02 and outer.extra == {"k": "v"} + assert tl.phases[0].extra == {"n": 3} + assert os.environ[it.ENV_DIR] == str(tmp_path / "run") + + +def test_worker_helpers_write_jsonl_and_ingest(tmp_path, monkeypatch): + run_dir = tmp_path / "run" + tl = it.InitTimeline.start(enabled=True, run_dir=run_dir) + # Simulate a worker: no in-process timeline, env set by the parent. + it.InitTimeline.reset() + monkeypatch.setenv(it.ENV_STAGE_ID, "1") + monkeypatch.setenv(it.ENV_REPLICA_ID, "2") + with it.worker_phase("load_weights", gib=1.5): + time.sleep(0.005) + it.worker_mark("proc_entry") + files = list(run_dir.glob("worker_*.jsonl")) + assert len(files) == 1 + rows = [json.loads(line) for line in files[0].read_text().splitlines()] + assert {r["name"] for r in rows} == {"load_weights", "proc_entry"} + assert all(r["stage_id"] == 1 and r["replica_id"] == 2 and r["role"] == "worker" for r in rows) + # Back in the engine process: ingest (twice -> idempotent). + it.InitTimeline._current = tl + assert tl.ingest_workers() == 2 + assert tl.ingest_workers() == 0 + assert [p.name for p in tl.phases] == sorted( + [p.name for p in tl.phases], key=lambda n: [p.start for p in tl.phases if p.name == n][0] + ) + + +def test_table_dump_and_suggested_timeouts(tmp_path): + tl = it.InitTimeline(enabled=True, run_dir=None, t0=1000.0) + tl.record(it.InitPhase("g2_spawn", 1000.0, 1010.0, stage_id=0, role="engine", pid=1)) + tl.record(it.InitPhase("load_weights", 1002.0, 1200.0, stage_id=0, role="worker", pid=2)) + tl.record(it.InitPhase("load_weights", 1210.0, 1260.0, stage_id=1, role="worker", pid=3)) + tl.t_end = 1500.0 + assert tl.total_s == pytest.approx(500.0) + table = tl.format_table() + assert "load_weights" in table and "g2_spawn" in table and "total=500.00s" in table + d = tl.to_dict() + assert d["schema"] == 1 and len(d["phases"]) == 3 + out = tmp_path / "tl.json" + tl.dump(out) + assert json.loads(out.read_text())["total_s"] == pytest.approx(500.0) + # stage 0 span = 1000..1200 = 200 s -> 300 (== default floor); total 500 s -> 750 + sug = tl.suggest_timeouts() + assert sug == {"stage_init_timeout": 300, "init_timeout": 750} + # short runs floor at the defaults + short = it.InitTimeline(enabled=True, t0=0.0) + short.record(it.InitPhase("x", 0.0, 1.0, stage_id=0)) + short.t_end = 2.0 + assert short.suggest_timeouts() == {"stage_init_timeout": 300, "init_timeout": 600} + + +def test_timed_phase_decorator_records_in_engine_process(tmp_path): + tl = it.InitTimeline.start(enabled=True, run_dir=tmp_path / "run") + + class Runner: + @it.timed_phase("capture") + def capture(self, x): + return x * 2 + + assert Runner().capture(21) == 42 + assert [p.name for p in tl.phases] == ["capture"] + assert tl.phases[0].role == "engine" diff --git a/vllm_omni/engine/init_timeline.py b/vllm_omni/engine/init_timeline.py new file mode 100644 index 000000000..c16064192 --- /dev/null +++ b/vllm_omni/engine/init_timeline.py @@ -0,0 +1,347 @@ +"""Initialization timeline: phase timers for multi-stage engine start-up. + +Two roles write into one timeline: + +* the **engine** process (``AsyncOmniEngine`` + ``StageRuntime``) records phases + directly through :meth:`InitTimeline.phase`; +* **worker/engine-core** subprocesses append JSON lines to + ``$VLLM_OMNI_INIT_TIMELINE_DIR/worker_.jsonl`` through :func:`worker_phase` + / :func:`worker_mark` (they inherit the env at spawn time, so no plumbing + through constructor arguments is needed). The engine ingests those files + once start-up finishes (or times out) with :meth:`InitTimeline.ingest_workers`. + +Enable with ``VLLM_OMNI_INIT_TIMELINE=/path/to/timeline.json`` (a JSON dump is +written there) or ``log_stats=True`` (table logged at INFO only). Disabled +timelines cost one attribute check per phase. +""" + +from __future__ import annotations + +import functools +import json +import math +import os +import tempfile +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +ENV_OUT = "VLLM_OMNI_INIT_TIMELINE" +ENV_DIR = "VLLM_OMNI_INIT_TIMELINE_DIR" +ENV_STAGE_ID = "VLLM_OMNI_STAGE_ID" +ENV_REPLICA_ID = "VLLM_OMNI_REPLICA_ID" +SCHEMA_VERSION = 1 + +DEFAULT_STAGE_INIT_TIMEOUT_S = 300 +DEFAULT_INIT_TIMEOUT_S = 600 + + +@dataclass +class InitPhase: + name: str + start: float + end: float + stage_id: int | None = None + replica_id: int = 0 + pid: int = 0 + role: str = "engine" + extra: dict[str, Any] = field(default_factory=dict) + + @property + def dur_s(self) -> float: + return max(0.0, self.end - self.start) + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["dur_s"] = self.dur_s + return d + + +def _env_int(name: str) -> int | None: + raw = os.environ.get(name) + if raw is None or raw == "": + return None + try: + return int(raw) + except ValueError: + return None + + +class InitTimeline: + """Process-local collector of :class:`InitPhase` records.""" + + _current: InitTimeline | None = None + _lock = threading.Lock() + + def __init__(self, enabled: bool, run_dir: Path | None = None, t0: float | None = None): + self.enabled = bool(enabled) + self.run_dir = Path(run_dir) if run_dir is not None else None + self.t0 = t0 if t0 is not None else time.time() + self.t_end: float | None = None + self.phases: list[InitPhase] = [] + self._phase_lock = threading.Lock() + + # ------------------------------------------------------------ lifecycle + @classmethod + def start(cls, enabled: bool, run_dir: Path | str | None = None) -> InitTimeline: + """Create the process-global timeline (idempotent per process).""" + with cls._lock: + run_path: Path | None = None + if enabled: + run_path = Path(run_dir) if run_dir else Path(tempfile.mkdtemp(prefix="omni_init_timeline_")) + run_path.mkdir(parents=True, exist_ok=True) + os.environ[ENV_DIR] = str(run_path) + tl = cls(enabled=enabled, run_dir=run_path) + cls._current = tl + return tl + + @classmethod + def current(cls) -> InitTimeline | None: + return cls._current + + @classmethod + def reset(cls) -> None: + with cls._lock: + cls._current = None + + def finish(self) -> None: + self.t_end = time.time() + + # ------------------------------------------------------------- recording + def record(self, phase: InitPhase) -> None: + if not self.enabled: + return + with self._phase_lock: + self.phases.append(phase) + + @contextmanager + def phase( + self, name: str, stage_id: int | None = None, replica_id: int = 0, **extra: Any + ) -> Iterator[dict[str, Any]]: + """Time a block; the yielded dict can receive extra fields.""" + info: dict[str, Any] = dict(extra) + if not self.enabled: + yield info + return + start = time.time() + try: + yield info + finally: + self.record( + InitPhase( + name=name, + start=start, + end=time.time(), + stage_id=stage_id, + replica_id=replica_id, + pid=os.getpid(), + role="engine", + extra=info, + ) + ) + + def mark(self, name: str, stage_id: int | None = None, replica_id: int = 0, **extra: Any) -> None: + now = time.time() + self.record( + InitPhase( + name=name, + start=now, + end=now, + stage_id=stage_id, + replica_id=replica_id, + pid=os.getpid(), + role="engine", + extra=dict(extra), + ) + ) + + # ------------------------------------------------------------- ingestion + def ingest_workers(self) -> int: + """Merge worker JSONL files from ``run_dir``; returns the number of phases added.""" + if not self.enabled or self.run_dir is None or not self.run_dir.exists(): + return 0 + added = 0 + seen = {(p.pid, p.name, p.start) for p in self.phases} + for path in sorted(self.run_dir.glob("worker_*.jsonl")): + try: + lines = path.read_text(encoding="utf-8").splitlines() + except OSError: + continue + for line in lines: + line = line.strip() + if not line: + continue + try: + d = json.loads(line) + phase = InitPhase( + name=str(d["name"]), + start=float(d["start"]), + end=float(d["end"]), + stage_id=d.get("stage_id"), + replica_id=int(d.get("replica_id", 0) or 0), + pid=int(d.get("pid", 0) or 0), + role=str(d.get("role", "worker")), + extra=dict(d.get("extra") or {}), + ) + except (KeyError, ValueError, TypeError): + continue + key = (phase.pid, phase.name, phase.start) + if key in seen: + continue + seen.add(key) + self.phases.append(phase) + added += 1 + self.phases.sort(key=lambda p: (p.start, p.end)) + return added + + # ------------------------------------------------------------- reporting + @property + def total_s(self) -> float: + end = self.t_end if self.t_end is not None else time.time() + return max(0.0, end - self.t0) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": SCHEMA_VERSION, + "t0": self.t0, + "t_end": self.t_end, + "total_s": self.total_s, + "phases": [p.to_dict() for p in sorted(self.phases, key=lambda p: (p.start, p.end))], + } + + def dump(self, path: str | Path) -> None: + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w", encoding="utf-8") as f: + json.dump(self.to_dict(), f, indent=1) + + def stage_spans(self) -> dict[int | None, tuple[float, float]]: + spans: dict[int | None, tuple[float, float]] = {} + for p in self.phases: + lo, hi = spans.get(p.stage_id, (p.start, p.end)) + spans[p.stage_id] = (min(lo, p.start), max(hi, p.end)) + return spans + + def format_table(self) -> str: + rows = sorted(self.phases, key=lambda p: (p.start, p.end)) + header = f"{'stage':>5} {'role':<7} {'pid':>8} {'phase':<28} {'start_s':>9} {'dur_s':>9}" + lines = [f"[init_timeline] total={self.total_s:.2f}s phases={len(rows)}", header, "-" * len(header)] + for p in rows: + stage = "-" if p.stage_id is None else str(p.stage_id) + lines.append( + f"{stage:>5} {p.role:<7} {p.pid:>8} {p.name[:28]:<28} {p.start - self.t0:>9.2f} {p.dur_s:>9.2f}" + ) + return "\n".join(lines) + + def suggest_timeouts(self, factor: float = 1.5) -> dict[str, int]: + """Timeouts (seconds) that would cover this run with ``factor`` headroom.""" + stage_max = 0.0 + for stage_id, (lo, hi) in self.stage_spans().items(): + if stage_id is None: + continue + stage_max = max(stage_max, hi - lo) + return { + "stage_init_timeout": max(DEFAULT_STAGE_INIT_TIMEOUT_S, int(math.ceil(stage_max * factor))), + "init_timeout": max(DEFAULT_INIT_TIMEOUT_S, int(math.ceil(self.total_s * factor))), + } + + +# ------------------------------------------------------------------ workers +def _worker_run_dir() -> Path | None: + raw = os.environ.get(ENV_DIR) + return Path(raw) if raw else None + + +def _append_worker_phase(phase: InitPhase) -> None: + run_dir = _worker_run_dir() + if run_dir is None: + return + try: + run_dir.mkdir(parents=True, exist_ok=True) + with open(run_dir / f"worker_{os.getpid()}.jsonl", "a", encoding="utf-8") as f: + f.write(json.dumps(phase.to_dict()) + "\n") + except OSError: # never let telemetry break init + pass + + +def _record_any(phase: InitPhase) -> None: + """Record into the in-process timeline if this is the engine process, else to JSONL.""" + tl = InitTimeline.current() + if tl is not None and tl.enabled: + tl.record(phase) + return + _append_worker_phase(phase) + + +def timeline_enabled() -> bool: + tl = InitTimeline.current() + if tl is not None: + return tl.enabled + return _worker_run_dir() is not None + + +@contextmanager +def worker_phase( + name: str, stage_id: int | None = None, replica_id: int | None = None, **extra: Any +) -> Iterator[dict[str, Any]]: + """Time a block in any process; no-op unless a timeline is active.""" + info: dict[str, Any] = dict(extra) + if not timeline_enabled(): + yield info + return + start = time.time() + try: + yield info + finally: + _record_any( + InitPhase( + name=name, + start=start, + end=time.time(), + stage_id=stage_id if stage_id is not None else _env_int(ENV_STAGE_ID), + replica_id=replica_id if replica_id is not None else (_env_int(ENV_REPLICA_ID) or 0), + pid=os.getpid(), + role="engine" if InitTimeline.current() is not None else "worker", + extra=info, + ) + ) + + +def worker_mark(name: str, **extra: Any) -> None: + if not timeline_enabled(): + return + now = time.time() + _record_any( + InitPhase( + name=name, + start=now, + end=now, + stage_id=_env_int(ENV_STAGE_ID), + replica_id=_env_int(ENV_REPLICA_ID) or 0, + pid=os.getpid(), + role="engine" if InitTimeline.current() is not None else "worker", + extra=dict(extra), + ) + ) + + +def timed_phase(name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator form of :func:`worker_phase` for methods on the init path.""" + + def deco(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with worker_phase(name): + return fn(*args, **kwargs) + + return wrapper + + return deco diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index e483e1bd4..d8544986b 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -151,6 +151,13 @@ class StageEngineCoreProc(EngineCoreProc): # Setting this env var allows the same graceful fallback to work. os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0)) + if omni_stage_id is not None: + # Lets init-timeline records from this process and its spawned + # workers carry the logical stage id. + os.environ["VLLM_OMNI_STAGE_ID"] = str(int(omni_stage_id)) + from vllm_omni.engine.init_timeline import worker_mark, worker_phase + + worker_mark("proc_entry") # Patch the decoder type so process_input_sockets (started # during __init__) decodes OmniEngineCoreRequest (which @@ -178,11 +185,12 @@ class StageEngineCoreProc(EngineCoreProc): if omni_parallel_stage_init: _install_phase_locks(kwargs, local_dp_rank) - engine_core = StageEngineCoreProc( - *args, - engine_index=dp_rank, - **kwargs, - ) + with worker_phase("engine_core_init"): + engine_core = StageEngineCoreProc( + *args, + engine_index=dp_rank, + **kwargs, + ) # Each subprocess corresponds to exactly one omni replica with # its own OmniMasterServer allocation, so the heartbeat client diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 93e482b42..33ed4e6ec 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -29,6 +29,7 @@ from vllm_omni.distributed.omni_coordinator import ( RandomBalancer, RoundRobinBalancer, ) +from vllm_omni.engine.init_timeline import worker_phase from vllm_omni.engine.messages import ( EngineQueueMessage, RegisterRemoteReplicaMessage, @@ -753,7 +754,10 @@ class StageRuntime: # READY handshake. if not self._parallel_stage_init: g3_start = time.perf_counter() - with self._scoped_spawn_device_env(physical_devices): + with ( + worker_phase("g3_device_lock", stage_id=plan.metadata.stage_id, replica_id=plan.replica_id), + self._scoped_spawn_device_env(physical_devices), + ): lock_fds = acquire_device_locks( plan.metadata.stage_id, plan.engine_args_dict, @@ -792,11 +796,13 @@ class StageRuntime: # and only needs to cover the spawn (__enter__) — children # inherit the env at spawn time; the READY-wait needs no env. g2_start = time.perf_counter() - with self._replica_launch_lock: - with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - resources = launch_cm.__enter__() + with worker_phase("g2_spawn", stage_id=plan.metadata.stage_id, replica_id=plan.replica_id): + with self._replica_launch_lock: + with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): + resources = launch_cm.__enter__() g2_spawned = time.perf_counter() - launch_cm.__exit__(None, None, None) + with worker_phase("g2_ready_wait", stage_id=plan.metadata.stage_id, replica_id=plan.replica_id): + launch_cm.__exit__(None, None, None) logger.debug( "[stage_init] Stage-%s G2 spawn(locked)=%.3fs, READY(unlocked)=%.3fs", plan.metadata.stage_id, @@ -805,10 +811,21 @@ class StageRuntime: ) else: g2_start = time.perf_counter() - with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - g2_locked = time.perf_counter() - with launch_cm as resources: - pass + with worker_phase( + "g2_launch_lock_wait", stage_id=plan.metadata.stage_id, replica_id=plan.replica_id + ) as lock_info: + self._replica_launch_lock.acquire() + lock_info["acquired"] = True + try: + with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): + g2_locked = time.perf_counter() + with worker_phase( + "g2_spawn_and_ready", stage_id=plan.metadata.stage_id, replica_id=plan.replica_id + ): + with launch_cm as resources: + pass + finally: + self._replica_launch_lock.release() logger.debug( "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn+READY=%.3fs", plan.metadata.stage_id, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 0aa6832d0..c7e0688d3 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -177,8 +177,17 @@ class OmniBase(PDDisaggregationMixin): "You should instead, pass the keyword arguments used to initialize the engine args " "directly to this object's initializer." ) - stage_init_timeout = kwargs.pop("stage_init_timeout", 300) - init_timeout = kwargs.pop("init_timeout", 600) + # Tri-state: ``None`` means "not specified". Named deploy profiles + # (e.g. ``deploy_profile="edge"``) initialize both stages on one device + # and compile/capture at start, so they get larger defaults; the + # init timeline (VLLM_OMNI_INIT_TIMELINE) reports what a run needed. + stage_init_timeout = kwargs.pop("stage_init_timeout", None) + init_timeout = kwargs.pop("init_timeout", None) + if kwargs.get("deploy_profile"): + stage_init_timeout = 900 if stage_init_timeout is None else stage_init_timeout + init_timeout = 1500 if init_timeout is None else init_timeout + stage_init_timeout = 300 if stage_init_timeout is None else stage_init_timeout + init_timeout = 600 if init_timeout is None else init_timeout log_stats = kwargs.pop("log_stats", False) self._enable_ar_profiler = kwargs.pop("enable_ar_profiler", False) # NOTE: read-only lookup — must NOT pop. Popping here drops the key diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 0308cda7b..b22afc027 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.layers.custom_op import CustomOp +from vllm_omni.engine.init_timeline import timed_phase from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -716,6 +717,7 @@ class CodePredictorWrapper(nn.Module): return self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) + @timed_phase("predictor_setup_compile") def _setup_compile(self) -> None: """Lazily set up torch.compile with optional device graph capture.""" if self._compiled_model_fwd is not None: @@ -842,6 +844,7 @@ class CodePredictorWrapper(nn.Module): allowed = set(all_seq_lens) return sorted(seq_len for seq_len in self._prefix_graph_seq_lens if seq_len in allowed) + @timed_phase("predictor_warmup") def _warmup_buckets(self) -> None: """Warmup power-of-2 batch-size buckets to front-load Inductor compilation.""" max_bsz = self._vllm_config.scheduler_config.max_num_seqs @@ -897,6 +900,7 @@ class CodePredictorWrapper(nn.Module): self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids) logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes) + @timed_phase("predictor_cudagraph") def _capture_cuda_graphs(self) -> None: """Capture a CUDA graph per bucket using vLLM's global graph pool.""" from vllm.platforms import current_platform diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 87df87889..c1388ef97 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -13,6 +13,7 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm_omni.engine.init_timeline import timed_phase from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp @@ -148,6 +149,7 @@ class Qwen3TTSCode2Wav(nn.Module): return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] return [ids] + @timed_phase("code2wav_cudagraph") def _maybe_enable_decoder_cudagraph( self, *, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d293f6312..dd362e865 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -45,6 +45,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output +from vllm_omni.engine.init_timeline import worker_phase from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import ( @@ -469,8 +470,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex return engine_output_type, downstream_req_ids def capture_model(self) -> int: - result = super().capture_model() - self._capture_talker_mtp_graphs() + with worker_phase("cudagraph_capture"): + result = super().capture_model() + with worker_phase("talker_mtp_capture"): + self._capture_talker_mtp_graphs() return result def shutdown(self) -> None: @@ -522,11 +525,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex # 5. Release all CUDA graphs unconditionally (upstream only does this # on ROCm; on CUDA the graphs are only freed by Python GC during # interpreter shutdown, which is too late to prevent memory spikes). - from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper - from vllm.compilation.cuda_graph import CUDAGraphWrapper + # Skipped on CPU, where no graph wrapper is ever instantiated. + if getattr(self.device, "type", "cuda") != "cpu": + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.compilation.cuda_graph import CUDAGraphWrapper - CUDAGraphWrapper.clear_all_graphs() - BreakableCUDAGraphWrapper.clear_all_graphs() + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() # 6. Delegate to upstream shutdown (model = None, KV caches, workspace). super().shutdown() diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 654cd4dc5..09d26d0ef 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -12,6 +12,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask +from vllm_omni.engine.init_timeline import timed_phase from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner @@ -31,6 +32,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): model_runner_cls = GPUARModelRunner @instrument(span_name="Init device") + @timed_phase("init_device") def init_device(self): if self.device_config.device_type in ("cuda", "musa"): # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index d6a8d6c8e..1a3eb5973 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -11,6 +11,7 @@ from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.engine.init_timeline import timed_phase from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner @@ -30,6 +31,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): model_runner_cls = GPUGenerationModelRunner @instrument(span_name="Init device") + @timed_phase("init_device") def init_device(self): if self.device_config.device_type in ("cuda", "musa"): # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index eca15c16e..560e728e6 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -34,6 +34,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +from vllm_omni.engine.init_timeline import timed_phase, worker_phase from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -179,8 +180,10 @@ class OmniGPUModelRunner(GPUModelRunner): ) @instrument(span_name="Loading (GPU)") + @timed_phase("load_model") def load_model(self, *args, **kwargs) -> None: - super().load_model(*args, **kwargs) + with worker_phase("load_weights"): + super().load_model(*args, **kwargs) model = getattr(self, "model", None) override_fn = None if bool(getattr(model, "supports_sampled_token_ids_cpu_override", False)): @@ -197,6 +200,7 @@ class OmniGPUModelRunner(GPUModelRunner): if getattr(self.model, "logitsprocs_need_output_token_ids", False): self.input_batch.logitsprocs_need_output_token_ids = True + @timed_phase("talker_mtp_init") def _init_talker_mtp(self) -> None: # TODO move this model specific logic to a separate class # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS. @@ -221,12 +225,14 @@ class OmniGPUModelRunner(GPUModelRunner): hidden_size = int( getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size") ) - max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size) + # max_cudagraph_capture_size is None when graphs are disabled (e.g. the CPU platform). + max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size or 0) self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32) self.talker_mtp_inputs_embeds = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) + @timed_phase("attn_prewarm") def _prewarm_attention_capture_workspaces(self) -> None: capture_sizes = getattr(self.compilation_config, "cudagraph_capture_sizes", None) if not capture_sizes: