[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>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
f804bbd679
commit
073d31b9cd
14 files changed
+885
-24
No files matched your search
@@ -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())
|
||||
@@ -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"],
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"
|
||||
@@ -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_<pid>.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
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user