diff --git a/benchmarks/tts/hw_emulation.py b/benchmarks/tts/hw_emulation.py new file mode 100644 index 000000000..24ba953e7 --- /dev/null +++ b/benchmarks/tts/hw_emulation.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Emulate edge hardware classes on a development host and benchmark each. + +Each case sets ``VLLM_OMNI_HW_PROFILE`` to a synthetic hardware profile (so +``deploy_profile="auto"`` derives budgets/threads/dtype for that class), plus +host-level levers that really constrain execution: + +* CPU cases: ``ATEN_CPU_CAPABILITY`` (torch ISA dispatch: ``avx2`` / + ``avx512`` / ``default``) and ``taskset`` core pinning; run in a CPU vLLM + environment with ``VLLM_TARGET_DEVICE=cpu``. +* CUDA cases (Jetson-like / small discrete GPU): the profile overrides GPU + memory so the derived KV budgets and eager switches apply; the card's real + memory is not capped (documented in the result). + +Writes one ``stream_latency_bench.py`` JSON per case plus a summary table under +``benchmarks/tts/edge_results/hw_emulation//``. + + # CPU cases (no GPU): + VLLM_TARGET_DEVICE=cpu python benchmarks/tts/hw_emulation.py \\ + --cases x86_cpu_avx2,x86_cpu_avx512,x86_cpu_amx,arm64_cpu_logic + # GPU cases, one GPU through the scheduler: + gpu run --gpus 1 --timeout 40m --note "hw emulation gpu" -- env HF_HUB_OFFLINE=1 \\ + python benchmarks/tts/hw_emulation.py --cases jetson_orin_8g,jetson_orin_32g,cuda_discrete_8g +""" + +from __future__ import annotations + +import argparse +import datetime as _dt +import json +import os +import shutil +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +BENCH_DIR = Path(__file__).resolve().parent +REPO_ROOT = BENCH_DIR.parent.parent +sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni.edge.hardware_probe import HardwareProfile, probe # noqa: E402 + +GiB = 2**30 + + +@dataclass +class Case: + name: str + profile: dict[str, Any] # fields overriding the live probe + env: dict[str, str] = field(default_factory=dict) + cpus: list[int] | None = None # taskset pinning + gpu: bool = False + note: str = "" + + +def _cpu_profile(flags: list[str], n: int, ram_gib: int, arch: str = "x86_64", model: str = "") -> dict[str, Any]: + return { + "arch": arch, + "cpu_model": model or f"emulated {arch}", + "cpu_flags": flags, + "n_logical": n, + "n_physical": n, + "clusters": [{"max_khz": 3000000, "cpus": list(range(n))}], + "ram_total_bytes": ram_gib * GiB, + "ram_available_bytes": int(ram_gib * 0.75) * GiB, + "accelerator": "none", + "supports_bf16": "avx512_bf16" in flags or "bf16" in flags, + "supports_fp16": True, + "source": "emulation", + } + + +def _gpu_profile(name: str, vram_gib: int, unified: bool, ram_gib: int, arch: str) -> dict[str, Any]: + return { + "arch": arch, + "cpu_model": f"emulated {name} host", + "cpu_flags": ["neon", "dotprod", "fp16_arith"] if arch == "aarch64" else ["avx2", "avx512"], + "n_logical": 8, + "n_physical": 8, + "clusters": [{"max_khz": 2200000, "cpus": list(range(8))}], + "ram_total_bytes": ram_gib * GiB, + "ram_available_bytes": int(ram_gib * 0.75) * GiB, + "accelerator": "cuda_unified" if unified else "cuda_discrete", + "gpu_name": name, + "gpu_mem_bytes": vram_gib * GiB, + "gpu_sm_count": 16, + "gpu_capability": [8, 7] if unified else [8, 6], + "supports_bf16": True, + "supports_fp16": True, + "source": "emulation", + } + + +CASES: dict[str, Case] = { + "x86_cpu_avx2": Case( + "x86_cpu_avx2", + _cpu_profile(["avx2", "fma", "f16c"], 4, 8), + env={"ATEN_CPU_CAPABILITY": "avx2"}, + cpus=[0, 1, 2, 3], + note="4 cores, AVX2 dispatch, 8 GiB profile -> float32", + ), + "x86_cpu_avx512": Case( + "x86_cpu_avx512", + _cpu_profile(["avx2", "avx512", "avx512_bf16"], 8, 16), + env={"ATEN_CPU_CAPABILITY": "avx512"}, + cpus=list(range(8)), + note="8 cores, AVX-512 (+bf16) dispatch, 16 GiB", + ), + "x86_cpu_amx": Case( + "x86_cpu_amx", + _cpu_profile(["avx2", "avx512", "avx512_bf16", "amx_bf16", "amx_tile"], 16, 32), + cpus=list(range(16)), + note="16 cores, AMX bf16, 32 GiB", + ), + "arm64_cpu_logic": Case( + "arm64_cpu_logic", + _cpu_profile(["neon", "dotprod", "fp16_arith", "i8mm"], 4, 8, arch="aarch64", model="emulated RK3588"), + cpus=[0, 1, 2, 3], + note="derivation/binding logic only: ISA is x86 underneath (labelled logic-only)", + ), + "jetson_orin_8g": Case( + "jetson_orin_8g", + _gpu_profile("Jetson Orin 8GB (emulated)", 8, True, 8, "aarch64"), + gpu=True, + note="unified 8 GiB -> tight budgets, eager", + ), + "jetson_orin_32g": Case( + "jetson_orin_32g", + _gpu_profile("Jetson Orin 32GB (emulated)", 32, True, 32, "aarch64"), + gpu=True, + note="unified 32 GiB", + ), + "cuda_discrete_8g": Case( + "cuda_discrete_8g", + _gpu_profile("small discrete GPU (emulated)", 8, False, 16, "x86_64"), + gpu=True, + note="8 GiB VRAM -> eager + small KV", + ), +} + + +def build_profile(case: Case) -> HardwareProfile: + base = probe(use_torch=case.gpu).to_dict() + base.update(case.profile) + return HardwareProfile.from_dict(base) + + +def run_case(case: Case, args: argparse.Namespace, out_dir: Path, passthrough: list[str]) -> dict[str, Any]: + case_dir = out_dir / case.name + case_dir.mkdir(parents=True, exist_ok=True) + prof = build_profile(case) + prof_path = case_dir / "hw_profile.json" + prof_path.write_text(prof.to_json()) + env = dict(os.environ) + env.update(case.env) + env["VLLM_OMNI_HW_PROFILE"] = str(prof_path) + env["VLLM_OMNI_INIT_TIMELINE"] = str(case_dir / "timeline.json") + if not case.gpu: + env.setdefault("VLLM_TARGET_DEVICE", "cpu") + cmd: list[str] = [] + if case.cpus and shutil.which("taskset"): + cmd += ["taskset", "-c", ",".join(str(c) for c in case.cpus)] + cmd += [ + args.python, + str(BENCH_DIR / "stream_latency_bench.py"), + "--model", + args.model, + "--deploy-profile", + "auto", + "--query-type", + "CustomVoice", + "--txt-prompts", + args.prompts, + "--warmup", + "1", + "--repeat", + str(args.repeat), + "--output-dir", + str(case_dir / "bench"), + "--tag", + case.name, + *passthrough, + ] + with open(case_dir / "stdout.log", "w") as log: + try: + proc = subprocess.run( + cmd, env=env, stdout=log, stderr=subprocess.STDOUT, timeout=args.timeout_s, cwd=str(REPO_ROOT) + ) + rc = proc.returncode + except subprocess.TimeoutExpired: + rc = -1 + row: dict[str, Any] = {"case": case.name, "note": case.note, "returncode": rc, "cmd": cmd, "hardware_class": None} + files = sorted((case_dir / "bench").rglob("*.json")) + if files: + doc = json.loads(files[-1].read_text()) + s = doc["summary"] + row.update( + init_s=doc["init"]["init_s"], + ttfa_p50=s["ttfa_ms"]["p50"], + playback_start_p50=s["playback_start_ms"]["p50"], + stall_p50=s["stall_at_ttfa_ms"]["p50"], + rtf_total_p50=s["rtf_total"]["p50"], + peak_rss_gib=doc["resources"]["peak_rss_bytes"] / GiB, + peak_gpu_mib=doc["resources"]["peak_gpu_mem_mib"], + ) + return row + + +def render_table(rows: list[dict[str, Any]]) -> str: + cols = [ + "case", + "rc", + "init_s", + "ttfa_p50", + "playback_start_p50", + "stall_p50", + "rtf_total_p50", + "peak_rss_gib", + "peak_gpu_mib", + "note", + ] + out = ["| " + " | ".join(cols) + " |", "|" + "|".join(["---"] * len(cols)) + "|"] + for r in rows: + vals = [r["case"], str(r.get("returncode"))] + for c in ("init_s", "ttfa_p50", "playback_start_p50", "stall_p50", "rtf_total_p50", "peak_rss_gib"): + v = r.get(c) + vals.append("" if v is None else (f"{v:.3f}" if c == "rtf_total_p50" else f"{v:.1f}")) + vals.append(str(r.get("peak_gpu_mib", ""))) + vals.append(r.get("note", "")) + out.append("| " + " | ".join(vals) + " |") + return "\n".join(out) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--cases", default="x86_cpu_avx2,x86_cpu_avx512,x86_cpu_amx,arm64_cpu_logic") + p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice") + p.add_argument("--prompts", default=str(BENCH_DIR / "prompts_12.txt")) + p.add_argument("--repeat", type=int, default=1) + p.add_argument("--timeout-s", type=int, default=2400) + p.add_argument("--output-dir", default=str(BENCH_DIR / "edge_results" / "hw_emulation")) + 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, c in CASES.items(): + print(f"{name}: gpu={c.gpu} env={c.env} cpus={c.cpus} {c.note}") + return 0 + names = [n.strip() for n in args.cases.split(",") if n.strip()] + unknown = [n for n in names if n not in CASES] + if unknown: + raise SystemExit(f"Unknown cases {unknown}; known: {sorted(CASES)}") + 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) + if "--init-timeout" not in passthrough: + passthrough += ["--init-timeout", "1500"] + if "--stage-init-timeout" not in passthrough: + passthrough += ["--stage-init-timeout", "1200"] + rows = [] + for name in names: + print(f"[hw_emulation] {name} ...", flush=True) + row = run_case(CASES[name], args, out_dir, passthrough) + rows.append(row) + print( + f"[hw_emulation] {name}: rc={row['returncode']} " + f"ttfa_p50={row.get('ttfa_p50')} rtf={row.get('rtf_total_p50')}", + flush=True, + ) + summary = {"model": args.model, "timestamp": ts, "rows": rows, "table": render_table(rows)} + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(render_table(rows)) + print(f"[hw_emulation] wrote {out_dir / 'summary.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/benchmarks/test_hw_emulation.py b/tests/benchmarks/test_hw_emulation.py new file mode 100644 index 000000000..019a8468f --- /dev/null +++ b/tests/benchmarks/test_hw_emulation.py @@ -0,0 +1,43 @@ +"""Tests for benchmarks/tts/hw_emulation.py case construction (CPU).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "benchmarks" / "tts")) +import hw_emulation as hw # noqa: E402 + +from vllm_omni.edge.hardware_probe import hardware_class # noqa: E402 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_cases_map_to_expected_classes(): + expected = { + "x86_cpu_avx2": "x86_cpu", + "x86_cpu_avx512": "x86_cpu", + "x86_cpu_amx": "x86_cpu", + "arm64_cpu_logic": "arm64_cpu", + "jetson_orin_8g": "jetson", + "jetson_orin_32g": "jetson", + "cuda_discrete_8g": "cuda_discrete", + } + for name, cls in expected.items(): + prof = hw.build_profile(hw.CASES[name]) + assert hardware_class(prof) == cls, name + assert prof.source == "emulation" + + +def test_avx2_case_disables_bf16_and_pins_cores(): + case = hw.CASES["x86_cpu_avx2"] + prof = hw.build_profile(case) + assert prof.supports_bf16 is False and prof.n_logical == 4 + assert case.env["ATEN_CPU_CAPABILITY"] == "avx2" and case.cpus == [0, 1, 2, 3] + + +def test_table_render_handles_missing_metrics(): + table = hw.render_table([{"case": "x", "returncode": -1, "note": "timeout"}]) + assert "| x | -1 |" in table and "timeout" in table diff --git a/tests/edge/__init__.py b/tests/edge/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/edge/test_adapt.py b/tests/edge/test_adapt.py new file mode 100644 index 000000000..7d06cfd45 --- /dev/null +++ b/tests/edge/test_adapt.py @@ -0,0 +1,109 @@ +"""Tests for vllm_omni.edge.adapt (CPU).""" + +from __future__ import annotations + +import pytest + +from vllm_omni.edge import adapt +from vllm_omni.edge.hardware_probe import HardwareProfile + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +GiB = 2**30 + + +def _x86(ram=32 * GiB, flags=("avx2", "avx512", "avx512_bf16", "amx_bf16"), n=16): + return HardwareProfile( + arch="x86_64", + cpu_flags=list(flags), + n_logical=n, + n_physical=n, + clusters=[{"max_khz": 3000000, "cpus": list(range(n))}], + ram_total_bytes=ram, + ram_available_bytes=ram - 2 * GiB, + supports_bf16="avx512_bf16" in flags, + ) + + +def test_dtype_rules(): + assert adapt.choose_dtype(_x86()) == "bfloat16" + assert adapt.choose_dtype(_x86(flags=("avx2",))) == "float32" + gpu = HardwareProfile(arch="x86_64", accelerator="cuda_discrete", gpu_mem_bytes=24 * GiB, supports_bf16=False) + assert adapt.choose_dtype(gpu) == "float16" + arm = HardwareProfile(arch="aarch64", cpu_flags=["neon", "fp16_arith"], supports_bf16=False) + assert adapt.choose_dtype(arm) == "float32" + + +def test_cpu_class_binds_disjoint_ranges_and_absolute_kv_budget(): + prof = _x86(ram=16 * GiB, n=8) + ov = adapt.derive_overrides(prof, weights_bytes=2 * GiB) + assert ov["hardware_class"] == "x86_cpu" + s0, s1 = ov["stages"][0], ov["stages"][1] + assert s0["devices"] == "cpu" and s1["devices"] == "cpu" + assert s0["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "0-4" + assert s1["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "5-7" + assert s0["enforce_eager"] is True and s1["enforce_eager"] is True + kv0, kv1 = s0["engine_args"]["kv_cache_memory_bytes"], s1["engine_args"]["kv_cache_memory_bytes"] + assert kv0 > 0 and kv1 > 0 and kv0 + kv1 == ov["memory_budget_bytes"] + assert kv0 + kv1 <= prof.ram_available_bytes + assert ov["connector_extra"]["codec_chunk_adaptive"] is True + assert ov["timeouts"]["stage_init_timeout"] >= 300 + + +def test_small_vram_gpu_forces_eager_and_tight_budgets(): + prof = HardwareProfile( + arch="x86_64", + accelerator="cuda_discrete", + gpu_name="small", + gpu_mem_bytes=6 * GiB, + supports_bf16=True, + ram_total_bytes=16 * GiB, + ) + ov = adapt.derive_overrides(prof, weights_bytes=3 * GiB) + assert ov["hardware_class"] == "cuda_discrete" + assert ov["stages"][0]["enforce_eager"] is True and ov["stages"][1]["enforce_eager"] is True + assert ov["stages"][0]["gpu_memory_utilization"] is None + total_kv = sum(ov["stages"][i]["engine_args"]["kv_cache_memory_bytes"] for i in (0, 1)) + assert total_kv <= 6 * GiB - 3 * GiB + assert "compilation_config" not in ov["stages"][0] + + +def test_jetson_unified_memory_uses_ram_and_graphs_off_stage1(): + prof = HardwareProfile( + arch="aarch64", accelerator="cuda_unified", gpu_mem_bytes=32 * GiB, ram_total_bytes=32 * GiB, supports_bf16=True + ) + ov = adapt.derive_overrides(prof, weights_bytes=4 * GiB) + assert ov["hardware_class"] == "jetson" + assert ov["stages"][1]["enforce_eager"] is True + assert ov["stages"][0]["enforce_eager"] is False + assert ov["stages"][0]["compilation_config"]["cudagraph_capture_sizes"] == [1, 2] + assert ov["stages"][0]["max_num_seqs"] == 2 + assert ov["timeouts"]["init_timeout"] >= ov["timeouts"]["stage_init_timeout"] + + +def test_budget_never_exceeds_pool_and_seqs_shrink(): + prof = _x86(ram=3 * GiB, n=4) + prof.ram_available_bytes = 3 * GiB + ov = adapt.derive_overrides(prof, weights_bytes=1 * GiB) + assert ov["memory_budget_bytes"] >= 0 + assert ov["stages"][0]["max_num_seqs"] >= 1 + + +def test_split_cpu_ranges_and_cli_flatten(): + assert adapt.split_cpu_ranges([0, 1, 2, 3, 4, 5]) == ("0-3", "4-5") + assert adapt.split_cpu_ranges([7]) == ("7", "7") + assert adapt.split_cpu_ranges([0, 2, 3, 4]) == ("0,2-3", "4") + assert adapt._ranges([0, 1, 2, 5, 7, 8]) == "0-2,5,7-8" + ov = adapt.derive_overrides(_x86(), weights_bytes=GiB) + flat = adapt.overrides_to_cli(ov) + assert flat["async_chunk"] is True and "0" in flat["stage_overrides"] and "stage_init_timeout" in flat + + +def test_estimate_weights_bytes(tmp_path): + (tmp_path / "a.safetensors").write_bytes(b"x" * 10) + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "b.safetensors").write_bytes(b"y" * 5) + (tmp_path / "c.bin").write_bytes(b"z" * 100) + assert adapt.estimate_weights_bytes(tmp_path) == 15 + assert adapt.estimate_weights_bytes(None) == 0 + assert adapt.estimate_weights_bytes(tmp_path / "missing") == 0 diff --git a/tests/edge/test_auto_profile.py b/tests/edge/test_auto_profile.py new file mode 100644 index 000000000..6a92ae75f --- /dev/null +++ b/tests/edge/test_auto_profile.py @@ -0,0 +1,105 @@ +"""Tests for deploy_profile='auto' materialization (CPU, no engine).""" + +from __future__ import annotations + +from unittest import mock + +import pytest +import yaml + +from vllm_omni.config.stage_config import _DEPLOY_DIR, load_deploy_config +from vllm_omni.edge import auto_profile, calibrate +from vllm_omni.edge.hardware_probe import HardwareProfile +from vllm_omni.entrypoints import utils as entry_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +GiB = 2**30 +EDGE = _DEPLOY_DIR / "edge" / "qwen3_tts.yaml" + + +def _arm_profile(): + return HardwareProfile( + arch="aarch64", + cpu_model="RK3588", + cpu_flags=["neon", "dotprod", "fp16_arith"], + n_logical=8, + n_physical=8, + clusters=[{"max_khz": 2400000, "cpus": [4, 5, 6, 7]}, {"max_khz": 1800000, "cpus": [0, 1, 2, 3]}], + ram_total_bytes=16 * GiB, + ram_available_bytes=12 * GiB, + supports_bf16=False, + ) + + +def test_materialize_arm64_cpu(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "cache")) + path, report = auto_profile.materialize_auto_deploy( + "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", EDGE, profile=_arm_profile(), out_dir=tmp_path + ) + assert report["hardware_class"] == "arm64_cpu" + doc = yaml.safe_load(open(path)) + assert "base_config" not in doc + s0, s1 = doc["stages"] + # class overlay pins eager + max_num_seqs 1; derivation adds cpu binding on the big cluster only + assert s0["enforce_eager"] is True and s1["enforce_eager"] is True + assert s0["max_num_seqs"] == 1 + assert s0["devices"] == "cpu" and s1["devices"] == "cpu" + assert s0["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "4-6" + assert s1["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "7" + assert s0["engine_args"]["dtype"] == "float32" + assert s0["engine_args"]["kv_cache_memory_bytes"] > 0 + extra = doc["connectors"]["connector_of_shared_memory"]["extra"] + assert extra["codec_chunk_adaptive"] is True and extra["codec_chunk_ramp"] == [2, 4, 8, 16, 25] + # The result is a valid deploy config. + cfg = load_deploy_config(path) + assert len(cfg.stages) == 2 and cfg.stages[0].engine_extras["dtype"] == "float32" + + +def test_calibration_cache_is_applied(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "cache")) + prof = _arm_profile() + fp = calibrate.hardware_fingerprint(prof, "m") + calibrate.save_calibration( + fp, + { + "codec_chunk_ramp": [3, 6, 12, 25], + "codec_chunk_min_frames": 3, + "codec_chunk_safety_margin_ms": 90.0, + "codec_chunk_adaptive": True, + }, + ) + path, report = auto_profile.materialize_auto_deploy("m", EDGE, profile=prof, out_dir=tmp_path) + doc = yaml.safe_load(open(path)) + extra = doc["connectors"]["connector_of_shared_memory"]["extra"] + assert extra["codec_chunk_ramp"] == [3, 6, 12, 25] and extra["codec_chunk_min_frames"] == 3 + assert report["calibration"]["codec_chunk_ramp"] == [3, 6, 12, 25] + + +def test_apply_deploy_profile_auto_uses_edge_base(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_OMNI_HW_PROFILE", str(tmp_path / "hw.json")) + (tmp_path / "hw.json").write_text(_arm_profile().to_json()) + with mock.patch.object(entry_utils, "resolve_model_config_path", return_value=str(_DEPLOY_DIR / "qwen3_tts.yaml")): + args = entry_utils.apply_deploy_profile("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", {"deploy_profile": "auto"}) + assert args["deploy_config"].endswith(".yaml") + doc = yaml.safe_load(open(args["deploy_config"])) + assert doc["stages"][0]["devices"] == "cpu" + assert doc["stages"][0]["max_num_batched_tokens"] == 512 # inherited from edge/qwen3_tts.yaml + + +def test_hardware_overlay_merge_semantics(): + doc = { + "stages": [{"stage_id": 0, "a": 1, "d": {"x": 1}}, {"stage_id": 1}], + "connectors": {"c": {"extra": {"k": 1}}}, + } + out = auto_profile.apply_derived_overrides( + doc, + { + "stages": {0: {"a": 2, "d": {"y": 2}, "gpu_memory_utilization": None}}, + "top": {"async_chunk": True}, + "connector_extra": {"z": 3}, + }, + ) + assert out["stages"][0] == {"stage_id": 0, "a": 2, "d": {"x": 1, "y": 2}} + assert out["async_chunk"] is True and out["connectors"]["c"]["extra"] == {"k": 1, "z": 3} + assert doc["stages"][0]["a"] == 1 # input untouched diff --git a/tests/edge/test_calibrate.py b/tests/edge/test_calibrate.py new file mode 100644 index 000000000..6db1f49c2 --- /dev/null +++ b/tests/edge/test_calibrate.py @@ -0,0 +1,75 @@ +"""Tests for vllm_omni.edge.calibrate (CPU).""" + +from __future__ import annotations + +import pytest + +from vllm_omni.edge import calibrate as cal +from vllm_omni.edge.hardware_probe import HardwareProfile + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SR = 24000 +FRAME = 1920 + + +def _chunks(rate_frames_per_s: float, sizes: list[int], ttfa_ms: float = 30.0): + """Synthesize arrivals for a generator producing ``rate`` frames/s.""" + chunks = [] + t = ttfa_ms + for i, n in enumerate(sizes): + if i > 0: + t += n * 1000.0 / rate_frames_per_s + chunks.append({"t_ms": t, "samples": n * FRAME}) + return chunks + + +def test_rate_estimate_from_measured_shape(): + # measured on L20X: 25 frames every ~160 ms -> ~156 frames/s + chunks = [ + {"t_ms": 30.0, "samples": FRAME}, + {"t_ms": 190.0, "samples": 25 * FRAME}, + {"t_ms": 350.0, "samples": 25 * FRAME}, + ] + rate = cal.frames_per_second_from_chunks(chunks, SR) + assert rate == pytest.approx(50 / 0.32, rel=1e-3) + assert cal.frames_per_second_from_chunks(chunks[:1], SR) is None + + +def test_fast_generator_keeps_small_first_chunk_and_no_adaptive(): + sched = cal.derive_chunk_schedule(_chunks(150.0, [1, 25, 25, 25]), SR) + assert sched["measured"] is True + assert sched["rtf"] < 0.2 + assert sched["codec_chunk_adaptive"] is False + assert sched["codec_chunk_min_frames"] == 1 + assert sched["codec_chunk_ramp"][-1] == 25 and sched["codec_chunk_ramp"][0] == 1 + + +def test_slow_generator_needs_lead_and_adaptive(): + # 10 frames/s < 12.5 real time: each 25-frame chunk takes 2.5 s to make but plays 2.0 s + sched = cal.derive_chunk_schedule(_chunks(10.0, [2, 4, 8, 16, 25, 25]), SR) + assert sched["rtf"] > 1.0 + assert sched["codec_chunk_adaptive"] is True + assert sched["codec_chunk_min_frames"] >= 7 # (40 + 20*25) / 80 -> 6.75 -> 7 + assert sched["codec_chunk_ramp"][-1] == 25 + assert all(a < b for a, b in zip(sched["codec_chunk_ramp"], sched["codec_chunk_ramp"][1:])) + extra = cal.connector_extra_from_calibration(sched) + assert extra["codec_chunk_adaptive"] is True and extra["codec_chunk_ramp"][-1] == 25 + + +def test_unmeasurable_falls_back_to_defaults(): + sched = cal.derive_chunk_schedule([{"t_ms": 30.0, "samples": FRAME}], SR) + assert sched["measured"] is False and sched["codec_chunk_ramp"] == [2, 4, 8, 16, 25] + + +def test_cache_roundtrip(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + prof = HardwareProfile(arch="x86_64", cpu_model="x", n_physical=8, gpu_name="", gpu_mem_bytes=0) + fp = cal.hardware_fingerprint(prof, "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice") + assert len(fp) == 16 + assert cal.load_calibration(fp) is None + p = cal.save_calibration(fp, {"codec_chunk_ramp": [1, 25]}) + assert p.parent == tmp_path / "omni_edge" + assert cal.load_calibration(fp) == {"codec_chunk_ramp": [1, 25]} + # different model -> different fingerprint + assert cal.hardware_fingerprint(prof, "other") != fp diff --git a/tests/edge/test_hardware_probe.py b/tests/edge/test_hardware_probe.py new file mode 100644 index 000000000..114eec882 --- /dev/null +++ b/tests/edge/test_hardware_probe.py @@ -0,0 +1,117 @@ +"""Tests for vllm_omni.edge.hardware_probe using synthetic /proc and /sys snapshots (CPU).""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from vllm_omni.edge import hardware_probe as hp + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +X86_CPUINFO = """processor\t: 0 +model name\t: Intel(R) Xeon(R) Platinum 8480C +flags\t\t: fpu sse avx2 avx512f avx512_bf16 amx_bf16 amx_tile fma f16c +processor\t: 1 +model name\t: Intel(R) Xeon(R) Platinum 8480C +flags\t\t: fpu sse avx2 avx512f avx512_bf16 amx_bf16 amx_tile fma f16c +""" + +ARM_CPUINFO = """processor\t: 0 +Features\t: fp asimd evtstrm aes asimdhp asimddp +processor\t: 1 +Features\t: fp asimd evtstrm aes asimdhp asimddp +processor\t: 2 +Features\t: fp asimd evtstrm aes asimdhp asimddp i8mm bf16 +processor\t: 3 +Features\t: fp asimd evtstrm aes asimdhp asimddp i8mm bf16 +Hardware\t: Rockchip RK3588 +""" + +MEMINFO = "MemTotal: 16000000 kB\nMemFree: 1000000 kB\nMemAvailable: 12000000 kB\n" + + +def _write_sysfs(root: Path, cpus: dict[int, tuple[int, int]], numa: int = 1, thermal: int = 2, tegra: bool = False): + for cpu, (max_khz, core_id) in cpus.items(): + d = root / "sys" / "devices" / "system" / "cpu" / f"cpu{cpu}" + (d / "cpufreq").mkdir(parents=True, exist_ok=True) + (d / "topology").mkdir(parents=True, exist_ok=True) + (d / "cpufreq" / "cpuinfo_max_freq").write_text(str(max_khz)) + (d / "topology" / "core_id").write_text(str(core_id)) + (d / "topology" / "physical_package_id").write_text("0") + for n in range(numa): + (root / "sys" / "devices" / "system" / "node" / f"node{n}").mkdir(parents=True, exist_ok=True) + for t in range(thermal): + (root / "sys" / "class" / "thermal" / f"thermal_zone{t}").mkdir(parents=True, exist_ok=True) + if tegra: + (root / "etc").mkdir(parents=True, exist_ok=True) + (root / "etc" / "nv_tegra_release").write_text("# R36 (release), REVISION: 4.0") + + +def _snapshot(tmp_path: Path, cpuinfo: str, meminfo: str = MEMINFO, **kw) -> Path: + root = tmp_path / "root" + (root / "proc").mkdir(parents=True) + (root / "proc" / "cpuinfo").write_text(cpuinfo) + (root / "proc" / "meminfo").write_text(meminfo) + _write_sysfs(root, **kw) + return root + + +def test_x86_amx_snapshot(tmp_path): + root = _snapshot(tmp_path, X86_CPUINFO, cpus={0: (3800000, 0), 1: (3800000, 0)}, numa=2) + p = hp.probe(root, use_torch=False) + assert p.arch == "x86_64" + assert p.has("avx2", "avx512", "avx512_bf16", "amx_bf16", "amx_tile") + assert p.n_logical == 2 and p.n_physical == 1 # two hyper-threads of one core + assert p.numa_nodes == 2 and p.thermal_zones == 2 + assert p.ram_total_bytes == 16000000 * 1024 and p.ram_available_bytes == 12000000 * 1024 + assert p.accelerator == "none" and p.supports_bf16 is True + assert hp.hardware_class(p) == hp.HW_CLASS_X86_CPU + assert p.clusters == [{"max_khz": 3800000, "cpus": [0, 1]}] + + +def test_arm_big_little_snapshot(tmp_path): + root = _snapshot(tmp_path, ARM_CPUINFO, cpus={0: (1800000, 0), 1: (1800000, 1), 2: (2400000, 2), 3: (2400000, 3)}) + p = hp.probe(root, use_torch=False) + assert p.arch == "aarch64" + assert p.has("neon", "fp16_arith", "dotprod", "i8mm", "bf16") + assert p.cpu_model == "Rockchip RK3588" + assert p.clusters[0] == {"max_khz": 2400000, "cpus": [2, 3]} # big cluster first + assert p.big_cpus == [2, 3] + assert p.n_physical == 4 + assert hp.hardware_class(p) == hp.HW_CLASS_ARM64_CPU + assert p.supports_bf16 is True and p.supports_fp16 is True + + +def test_tegra_snapshot_is_jetson_with_unified_memory(tmp_path): + root = _snapshot(tmp_path, ARM_CPUINFO, cpus={0: (2200000, 0)}, tegra=True) + p = hp.probe(root, use_torch=False) + assert p.accelerator == "cuda_unified" + assert p.gpu_mem_bytes == p.ram_total_bytes + assert hp.hardware_class(p) == hp.HW_CLASS_JETSON + + +def test_override_env_and_roundtrip(tmp_path, monkeypatch): + prof = hp.HardwareProfile( + arch="aarch64", cpu_flags=["neon"], n_logical=8, accelerator="npu:qualcomm", ram_total_bytes=8 * 2**30 + ) + path = tmp_path / "hw.json" + path.write_text(prof.to_json()) + monkeypatch.setenv(hp.ENV_OVERRIDE, str(path)) + loaded = hp.load_profile(use_torch=False) + assert loaded.source.startswith("override:") + assert loaded.accelerator == "npu:qualcomm" + assert hp.hardware_class(loaded) == hp.HW_CLASS_NPU_PHONE + # unknown keys are ignored on load + d = json.loads(prof.to_json()) + d["future_field"] = 1 + assert hp.HardwareProfile.from_dict(d).n_logical == 8 + assert "class=npu_phone" in hp.describe(loaded) + + +def test_live_probe_runs_on_this_host(): + p = hp.probe(use_torch=False) + assert p.n_logical >= 1 and p.ram_total_bytes > 0 + assert hp.hardware_class(p) in hp.HW_CLASSES diff --git a/vllm_omni/deploy/edge/hardware/README.md b/vllm_omni/deploy/edge/hardware/README.md new file mode 100644 index 000000000..2cf60e0f1 --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/README.md @@ -0,0 +1,22 @@ +# Hardware-class overlays for `deploy_profile="auto"` + +When a deploy profile is `auto`, vLLM-Omni probes the host +(`vllm_omni/edge/hardware_probe.py`), maps it to one of the classes below, and +materializes a deploy YAML as: `edge/.yaml` (or the model default) +→ overlay `edge/hardware/.yaml` (this directory; same `stages:` / +`connectors:` shape as a deploy file, merged per stage id) → derived overrides +from `vllm_omni/edge/adapt.py` (dtype, KV budgets, thread binding, eager +switches, adaptive chunking, timeouts) → cached calibration +(`vllm_omni/edge/calibrate.py`) if one exists for this hardware fingerprint. + +Class files are optional; add one to pin values that the derivation should not +touch for that class. Emulate a class on any host with +`VLLM_OMNI_HW_PROFILE=` (see `benchmarks/tts/hw_emulation.py`). + +| class | selected when | +|---|---| +| `jetson` | CUDA on unified memory (Tegra release file present) | +| `cuda_discrete` | CUDA device with its own memory | +| `arm64_cpu` | aarch64, no accelerator | +| `x86_cpu` | x86_64, no accelerator | +| `npu_phone` | `accelerator: npu:` (only via profile override; vLLM-Omni does not run on the NPU itself) | diff --git a/vllm_omni/deploy/edge/hardware/arm64_cpu.yaml b/vllm_omni/deploy/edge/hardware/arm64_cpu.yaml new file mode 100644 index 000000000..f2d6a2dcf --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/arm64_cpu.yaml @@ -0,0 +1,10 @@ +# ARM64 CPU boards: single utterance at a time, eager, generous chunk lead. +stages: + - stage_id: 0 + max_num_seqs: 1 + enforce_eager: true + async_scheduling: false + - stage_id: 1 + max_num_seqs: 1 + enforce_eager: true + async_scheduling: false diff --git a/vllm_omni/deploy/edge/hardware/cuda_discrete.yaml b/vllm_omni/deploy/edge/hardware/cuda_discrete.yaml new file mode 100644 index 000000000..56cc34ecc --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/cuda_discrete.yaml @@ -0,0 +1,3 @@ +# Discrete CUDA GPU (desktop / small server): defaults from edge/.yaml apply; +# nothing pinned here so the measured init_profile results decide graphs vs eager. +stages: [] diff --git a/vllm_omni/deploy/edge/hardware/jetson.yaml b/vllm_omni/deploy/edge/hardware/jetson.yaml new file mode 100644 index 000000000..1073e254e --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/jetson.yaml @@ -0,0 +1,8 @@ +# Jetson (unified memory): keep the decoder eager, small batches; budgets are +# absolute bytes derived at runtime, so gpu_memory_utilization is not used. +stages: + - stage_id: 0 + max_num_seqs: 2 + - stage_id: 1 + max_num_seqs: 2 + enforce_eager: true diff --git a/vllm_omni/deploy/edge/hardware/npu_phone.yaml b/vllm_omni/deploy/edge/hardware/npu_phone.yaml new file mode 100644 index 000000000..84b281839 --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/npu_phone.yaml @@ -0,0 +1,3 @@ +# Phone NPUs do not run vLLM-Omni; this class is only meaningful for the +# codec-token client (decoder backend selection). Kept for completeness. +stages: [] diff --git a/vllm_omni/deploy/edge/hardware/x86_cpu.yaml b/vllm_omni/deploy/edge/hardware/x86_cpu.yaml new file mode 100644 index 000000000..fc947f5d5 --- /dev/null +++ b/vllm_omni/deploy/edge/hardware/x86_cpu.yaml @@ -0,0 +1,10 @@ +# x86 mini-PC / laptop CPU: eager, 1-2 utterances, thread ranges derived at runtime. +stages: + - stage_id: 0 + max_num_seqs: 2 + enforce_eager: true + async_scheduling: false + - stage_id: 1 + max_num_seqs: 2 + enforce_eager: true + async_scheduling: false diff --git a/vllm_omni/edge/__init__.py b/vllm_omni/edge/__init__.py new file mode 100644 index 000000000..b60ede0c6 --- /dev/null +++ b/vllm_omni/edge/__init__.py @@ -0,0 +1 @@ +"""Edge-deployment helpers: hardware probing, config derivation, calibration.""" diff --git a/vllm_omni/edge/adapt.py b/vllm_omni/edge/adapt.py new file mode 100644 index 000000000..0fc24c06e --- /dev/null +++ b/vllm_omni/edge/adapt.py @@ -0,0 +1,199 @@ +"""Derive per-stage deploy overrides from a :class:`HardwareProfile`. + +The output is a plain dict in the deploy-YAML shape (``{"stages": {idx: {...}}, +"top": {...}, "connector_extra": {...}, "timeouts": {...}}``) so it can be +merged as CLI overrides or written as a variant YAML. Rules are deliberately +simple and documented inline; measured runs (``benchmarks/tts/hw_emulation.py``) +are the way to refine the constants in ``CLASS_TABLE``. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from vllm_omni.edge.hardware_probe import ( + HW_CLASS_ARM64_CPU, + HW_CLASS_CUDA_DISCRETE, + HW_CLASS_JETSON, + HW_CLASS_NPU_PHONE, + HW_CLASS_X86_CPU, + HardwareProfile, + hardware_class, +) + +GiB = 2**30 + + +@dataclass(frozen=True) +class ClassDefaults: + max_num_seqs: int + max_model_len_stage0: int + max_model_len_stage1: int + headroom_bytes: int # left free for the OS / other apps + kv_fraction_stage0: float # share of the KV budget for the AR stage + chunk_min_frames: int + chunk_safety_margin_ms: float + enforce_eager_stage0: bool + enforce_eager_stage1: bool + small_vram_bytes: int # below this, graphs off and budgets tight + read_bw_bytes_per_s: float # for the load-time cost model + + +CLASS_TABLE: dict[str, ClassDefaults] = { + HW_CLASS_CUDA_DISCRETE: ClassDefaults(4, 4096, 65536, 2 * GiB, 0.7, 2, 50.0, False, False, 8 * GiB, 2.0e9), + HW_CLASS_JETSON: ClassDefaults(2, 2048, 32768, 3 * GiB, 0.7, 2, 80.0, False, True, 16 * GiB, 0.8e9), + HW_CLASS_X86_CPU: ClassDefaults(2, 2048, 16384, 2 * GiB, 0.7, 2, 120.0, True, True, 0, 1.0e9), + HW_CLASS_ARM64_CPU: ClassDefaults(1, 2048, 16384, 1 * GiB, 0.7, 3, 160.0, True, True, 0, 0.4e9), + HW_CLASS_NPU_PHONE: ClassDefaults(1, 2048, 16384, 1 * GiB, 0.7, 4, 200.0, True, True, 0, 0.3e9), +} + + +def estimate_weights_bytes(model_dir: str | Path | None) -> int: + """Sum of ``*.safetensors`` under ``model_dir`` (0 when unknown).""" + if not model_dir: + return 0 + root = Path(model_dir) + if not root.exists(): + return 0 + return sum(p.stat().st_size for p in root.rglob("*.safetensors")) + + +def choose_dtype(profile: HardwareProfile) -> str: + if profile.accelerator.startswith("cuda"): + return "bfloat16" if profile.supports_bf16 else "float16" + if profile.supports_bf16: + return "bfloat16" + # CPUs without bf16: fp16 matmuls are slow/unsupported on most CPUs -> fp32. + return "float32" + + +def memory_budget_bytes(profile: HardwareProfile, weights_bytes: int, defaults: ClassDefaults) -> int: + """Bytes available for KV caches after weights and headroom.""" + if profile.accelerator == "cuda_discrete" and profile.gpu_mem_bytes: + pool = profile.gpu_mem_bytes + elif profile.accelerator == "cuda_unified": + pool = profile.gpu_mem_bytes or profile.ram_total_bytes + else: + pool = profile.ram_available_bytes or profile.ram_total_bytes + # Activations + allocator fragmentation: 15 % of weights, at least 512 MiB. + activations = max(int(0.15 * weights_bytes), 512 * 2**20) + return max(0, pool - weights_bytes - activations - defaults.headroom_bytes) + + +def split_cpu_ranges(cpus: list[int], share_stage0: float = 2 / 3) -> tuple[str, str]: + """Split a CPU list into two disjoint range strings for stage 0 / stage 1.""" + if not cpus: + return "", "" + n0 = max(1, int(round(len(cpus) * share_stage0))) + if len(cpus) > 1: + n0 = min(n0, len(cpus) - 1) + a, b = cpus[:n0], cpus[n0:] or cpus[-1:] + return _ranges(a), _ranges(b) + + +def _ranges(cpus: list[int]) -> str: + cpus = sorted(set(cpus)) + parts: list[str] = [] + start = prev = cpus[0] + for c in cpus[1:]: + if c == prev + 1: + prev = c + continue + parts.append(f"{start}-{prev}" if start != prev else f"{start}") + start = prev = c + parts.append(f"{start}-{prev}" if start != prev else f"{start}") + return ",".join(parts) + + +def derive_overrides( + profile: HardwareProfile, + *, + weights_bytes: int = 0, + n_stages: int = 2, + max_cpus: int | None = None, +) -> dict[str, Any]: + """Derive deploy overrides for an AR-talker + generation-decoder pipeline.""" + cls = hardware_class(profile) + d = CLASS_TABLE[cls] + dtype = choose_dtype(profile) + budget = memory_budget_bytes(profile, weights_bytes, d) + kv0 = int(budget * d.kv_fraction_stage0) + kv1 = int(budget - kv0) + # Never let KV budgets allow preemption: require room for max_num_seqs x max_model_len + # at a conservative 8 KiB/token for the talker (measured later, see WP8 tests). + max_num_seqs = d.max_num_seqs + per_seq_bytes = d.max_model_len_stage0 * 8 * 1024 + while max_num_seqs > 1 and max_num_seqs * per_seq_bytes > max(kv0, 1): + max_num_seqs -= 1 + + small_vram = ( + profile.accelerator.startswith("cuda") + and (profile.gpu_mem_bytes or profile.ram_total_bytes) < d.small_vram_bytes + ) + eager0 = d.enforce_eager_stage0 or small_vram + eager1 = d.enforce_eager_stage1 or small_vram + + stage0: dict[str, Any] = { + "max_num_seqs": max_num_seqs, + "max_model_len": d.max_model_len_stage0, + "enforce_eager": eager0, + "engine_args": {"dtype": dtype, "kv_cache_memory_bytes": kv0}, + } + stage1: dict[str, Any] = { + "max_num_seqs": max_num_seqs, + "max_model_len": d.max_model_len_stage1, + "enforce_eager": eager1, + "engine_args": {"dtype": dtype, "kv_cache_memory_bytes": kv1}, + } + if profile.accelerator.startswith("cuda"): + stage0["gpu_memory_utilization"] = None # absolute budget replaces the fraction + stage1["gpu_memory_utilization"] = None + if not eager0: + stage0["compilation_config"] = {"cudagraph_capture_sizes": sorted({1, 2, max_num_seqs})} + else: + cpus = profile.big_cpus + if max_cpus: + cpus = cpus[:max_cpus] + r0, r1 = split_cpu_ranges(cpus) + stage0["devices"] = "cpu" + stage1["devices"] = "cpu" + stage0["env"] = {"VLLM_CPU_OMP_THREADS_BIND": r0} + stage1["env"] = {"VLLM_CPU_OMP_THREADS_BIND": r1} + + connector_extra = { + "codec_chunk_adaptive": True, + "codec_chunk_min_frames": d.chunk_min_frames, + "codec_chunk_safety_margin_ms": d.chunk_safety_margin_ms, + "decode_batch_max_size": 1, + } + load_s = (weights_bytes / d.read_bw_bytes_per_s) if weights_bytes else 60.0 + compile_s = 0.0 if (eager0 and eager1) else 240.0 + stage_init = int(math.ceil((load_s + compile_s) * 2.0 + 60)) + timeouts = {"stage_init_timeout": max(300, stage_init), "init_timeout": max(600, stage_init * max(1, n_stages))} + return { + "hardware_class": cls, + "dtype": dtype, + "memory_budget_bytes": budget, + "stages": {0: stage0, 1: stage1}, + "top": {"async_chunk": True}, + "connector_extra": connector_extra, + "timeouts": timeouts, + "notes": [ + f"class={cls} dtype={dtype} budget={budget / GiB:.2f}GiB weights={weights_bytes / GiB:.2f}GiB", + f"small_vram={small_vram} eager=({eager0},{eager1}) max_num_seqs={max_num_seqs}", + ], + } + + +def overrides_to_cli(overrides: dict[str, Any]) -> dict[str, Any]: + """Flatten for ``StageConfigFactory.create_from_model(cli_overrides=...)``-style consumers.""" + flat: dict[str, Any] = {} + for k, v in overrides.get("top", {}).items(): + flat[k] = v + flat["stage_overrides"] = {str(i): s for i, s in overrides.get("stages", {}).items()} + flat["connector_extra"] = dict(overrides.get("connector_extra", {})) + flat.update({k: v for k, v in overrides.get("timeouts", {}).items()}) + return flat diff --git a/vllm_omni/edge/auto_profile.py b/vllm_omni/edge/auto_profile.py new file mode 100644 index 000000000..1ed41cf33 --- /dev/null +++ b/vllm_omni/edge/auto_profile.py @@ -0,0 +1,160 @@ +"""``deploy_profile="auto"``: materialize a hardware-adapted deploy YAML. + +Chain: base deploy YAML (``edge/.yaml`` if present, else the model +default) -> hardware-class overlay (``edge/hardware/.yaml``) -> derived +overrides (:func:`vllm_omni.edge.adapt.derive_overrides`) -> cached calibration +(:mod:`vllm_omni.edge.calibrate`). The result is written to a temp file and +returned as a normal ``deploy_config`` path, so nothing downstream changes. +""" + +from __future__ import annotations + +import atexit +import copy +import os +import tempfile +from pathlib import Path +from typing import Any + +import yaml +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import _DEPLOY_DIR, resolve_deploy_yaml +from vllm_omni.edge import calibrate +from vllm_omni.edge.adapt import derive_overrides, estimate_weights_bytes +from vllm_omni.edge.hardware_probe import HardwareProfile, describe, hardware_class, load_profile + +logger = init_logger(__name__) + +HARDWARE_DIR = _DEPLOY_DIR / "edge" / "hardware" + + +def _merge_stage_list(base_stages: list[dict[str, Any]], overlay_stages: list[dict[str, Any]]) -> list[dict[str, Any]]: + by_id = {int(s.get("stage_id", i)): dict(s) for i, s in enumerate(base_stages)} + for s in overlay_stages or []: + sid = int(s.get("stage_id", 0)) + target = by_id.setdefault(sid, {"stage_id": sid}) + for k, v in s.items(): + if isinstance(v, dict) and isinstance(target.get(k), dict): + target[k] = {**target[k], **v} + else: + target[k] = v + return [by_id[k] for k in sorted(by_id)] + + +def _connector_extra(doc: dict[str, Any]) -> dict[str, Any] | None: + connectors = doc.get("connectors") + if not isinstance(connectors, dict) or not connectors: + return None + first = next(iter(connectors.values())) + if not isinstance(first, dict): + return None + return first.setdefault("extra", {}) + + +def apply_hardware_overlay(doc: dict[str, Any], hw_class: str) -> dict[str, Any]: + path = HARDWARE_DIR / f"{hw_class}.yaml" + if not path.exists(): + return doc + with open(path, encoding="utf-8") as f: + overlay = yaml.safe_load(f) or {} + out = copy.deepcopy(doc) + if overlay.get("stages"): + out["stages"] = _merge_stage_list(out.get("stages", []), overlay["stages"]) + if isinstance(overlay.get("connectors"), dict): + extra = _connector_extra(out) + ov_extra = _connector_extra(overlay) + if extra is not None and ov_extra: + extra.update(ov_extra) + for k, v in overlay.items(): + if k not in ("stages", "connectors", "platforms", "base_config"): + out[k] = v + return out + + +def apply_derived_overrides(doc: dict[str, Any], overrides: dict[str, Any]) -> dict[str, Any]: + out = copy.deepcopy(doc) + stage_updates = [] + for sid, upd in overrides.get("stages", {}).items(): + entry: dict[str, Any] = {"stage_id": int(sid)} + for k, v in upd.items(): + if k == "gpu_memory_utilization" and v is None: + continue # keep the YAML fraction unless an absolute budget replaces it + entry[k] = v + stage_updates.append(entry) + out["stages"] = _merge_stage_list(out.get("stages", []), stage_updates) + for k, v in overrides.get("top", {}).items(): + out[k] = v + extra = _connector_extra(out) + if extra is not None: + extra.update(overrides.get("connector_extra", {})) + return out + + +def apply_calibration(doc: dict[str, Any], cal: dict[str, Any] | None) -> dict[str, Any]: + if not cal: + return doc + out = copy.deepcopy(doc) + extra = _connector_extra(out) + if extra is not None: + extra.update(calibrate.connector_extra_from_calibration(cal)) + return out + + +def _model_dir(model: str) -> str | None: + if os.path.isdir(model): + return model + try: + from huggingface_hub import snapshot_download + + return snapshot_download(model, local_files_only=True) + except Exception: + return None + + +def materialize_auto_deploy( + model: str, + base_yaml: str | Path, + *, + profile: HardwareProfile | None = None, + out_dir: str | Path | None = None, + use_calibration: bool = True, +) -> tuple[str, dict[str, Any]]: + """Build the adapted deploy YAML; returns ``(path, report)``.""" + prof = profile or load_profile() + hw_class = hardware_class(prof) + doc = resolve_deploy_yaml(Path(base_yaml)) + doc.pop("base_config", None) + # ``platforms:`` overlays are applied later by the config factory and would + # override the derived values (thread binding, budgets, eager switches); + # the hardware-class overlay + derivation supersede them in auto mode. + doc.pop("platforms", None) + doc = apply_hardware_overlay(doc, hw_class) + weights = estimate_weights_bytes(_model_dir(model)) + overrides = derive_overrides(prof, weights_bytes=weights, n_stages=len(doc.get("stages", [])) or 2) + doc = apply_derived_overrides(doc, overrides) + cal = None + if use_calibration: + cal = calibrate.load_calibration(calibrate.hardware_fingerprint(prof, model)) + doc = apply_calibration(doc, cal) + target_dir = Path(out_dir) if out_dir else None + if target_dir is not None: + target_dir.mkdir(parents=True, exist_ok=True) + fd, path = tempfile.mkstemp(prefix="auto_", suffix=".yaml", dir=str(target_dir)) + else: + fd, path = tempfile.mkstemp(prefix=f"omni_auto_{hw_class}_", suffix=".yaml") + atexit.register(lambda p=path: Path(p).unlink(missing_ok=True)) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(f"# auto-generated by vllm_omni.edge.auto_profile for {model}\n# {describe(prof)}\n") + yaml.safe_dump(doc, f, sort_keys=False) + report = { + "hardware_class": hw_class, + "profile": prof.to_dict(), + "base_yaml": str(base_yaml), + "weights_bytes": weights, + "overrides": overrides, + "calibration": cal, + "deploy_config": path, + } + logger.info("[auto_profile] %s -> %s (%s)", hw_class, path, "; ".join(overrides.get("notes", []))) + return path, report diff --git a/vllm_omni/edge/calibrate.py b/vllm_omni/edge/calibrate.py new file mode 100644 index 000000000..94b57884f --- /dev/null +++ b/vllm_omni/edge/calibrate.py @@ -0,0 +1,144 @@ +"""Calibration of the codec chunk schedule from measured chunk timings. + +Given the chunk arrivals of a warm-up request (as recorded by +``benchmarks/tts/stream_latency_bench.py``), derive parameters for the runtime +adaptive chunk controller (``codec_chunk_adaptive`` / ``codec_chunk_min_frames`` +/ ``codec_chunk_safety_margin_ms``) and a static ramp that would have played +back without underruns. Results are cached per hardware fingerprint so the +next start can apply them without re-measuring. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +from vllm_omni.edge.hardware_probe import HardwareProfile, hardware_class + +FRAME_MS = 80.0 # 12.5 Hz codec + + +def hardware_fingerprint(profile: HardwareProfile, model: str) -> str: + key = json.dumps( + { + "class": hardware_class(profile), + "cpu": profile.cpu_model, + "flags": sorted(profile.cpu_flags), + "gpu": profile.gpu_name, + "gpu_mem": profile.gpu_mem_bytes // (256 * 2**20), + "cores": profile.n_physical, + "model": model, + }, + sort_keys=True, + ) + return hashlib.sha1(key.encode()).hexdigest()[:16] + + +def cache_dir() -> Path: + root = os.environ.get("VLLM_CACHE_ROOT") or os.path.join(os.path.expanduser("~"), ".cache", "vllm") + return Path(root) / "omni_edge" + + +def cache_path(fingerprint: str) -> Path: + return cache_dir() / f"{fingerprint}.json" + + +def load_calibration(fingerprint: str) -> dict[str, Any] | None: + p = cache_path(fingerprint) + if not p.exists(): + return None + try: + return json.loads(p.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def save_calibration(fingerprint: str, data: dict[str, Any]) -> Path: + p = cache_path(fingerprint) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(data, indent=1), encoding="utf-8") + return p + + +def frames_per_second_from_chunks(chunks: list[dict[str, Any]], sr: int, frame_ms: float = FRAME_MS) -> float | None: + """Steady-state generation rate (codec frames per wall second) from chunk arrivals.""" + audio = [c for c in chunks if int(c.get("samples", 0)) > 0] + if len(audio) < 2: + return None + first, last = audio[0], audio[-1] + frames = sum(int(c["samples"]) for c in audio[1:]) / sr * 1000.0 / frame_ms + wall_s = (float(last["t_ms"]) - float(first["t_ms"])) / 1000.0 + if wall_s <= 0: + return None + return frames / wall_s + + +def derive_chunk_schedule( + chunks: list[dict[str, Any]], + sr: int, + *, + steady_frames: int = 25, + frame_ms: float = FRAME_MS, + jitter_ms: float = 40.0, +) -> dict[str, Any]: + """Derive adaptive-controller parameters and a static ramp from one measured request. + + Rule: the audio buffered before waiting for chunk ``i+1`` must cover the + time to produce it. With generation rate ``r`` frames/s (real-time factor + ``rtf = 12.5 / r``), producing ``n`` frames takes ``n / r`` seconds while + ``n`` frames play for ``n * 0.08`` seconds, so a chunk of size ``n`` is safe + once ``n * 0.08 >= n / r + margin``; the first chunk therefore sets the + playback lead. We choose ``min_frames`` so the lead covers ``margin`` and + build a doubling ramp up to ``steady_frames``. + """ + rate = frames_per_second_from_chunks(chunks, sr, frame_ms) + audio = [c for c in chunks if int(c.get("samples", 0)) > 0] + ttfa_ms = float(audio[0]["t_ms"]) if audio else None + if rate is None or rate <= 0: + return { + "measured": False, + "codec_chunk_adaptive": True, + "codec_chunk_min_frames": 2, + "codec_chunk_safety_margin_ms": 50.0, + "codec_chunk_ramp": [2, 4, 8, 16, steady_frames], + "ttfa_ms": ttfa_ms, + } + rtf = (1000.0 / frame_ms) / rate # wall seconds per audio second + # Time to produce one frame vs its playback duration. + produce_ms_per_frame = 1000.0 / rate + lead_needed_ms = jitter_ms + max(0.0, produce_ms_per_frame - frame_ms) * steady_frames + min_frames = max(1, int(-(-lead_needed_ms // frame_ms))) # ceil + min_frames = min(min_frames, steady_frames) + margin_ms = max(20.0, min(400.0, jitter_ms + produce_ms_per_frame)) + ramp: list[int] = [] + n = min_frames + while n < steady_frames: + ramp.append(n) + n *= 2 + ramp.append(steady_frames) + if len(ramp) < 2: + ramp = [max(1, steady_frames // 2), steady_frames] + return { + "measured": True, + "rate_frames_per_s": rate, + "rtf": rtf, + "ttfa_ms": ttfa_ms, + "codec_chunk_adaptive": rtf > 0.5, # only worth it when generation is not far faster than real time + "codec_chunk_min_frames": min_frames, + "codec_chunk_safety_margin_ms": margin_ms, + "codec_chunk_ramp": ramp, + } + + +def connector_extra_from_calibration(cal: dict[str, Any]) -> dict[str, Any]: + out = { + "codec_chunk_ramp": list(cal.get("codec_chunk_ramp", [2, 4, 8, 16, 25])), + "codec_chunk_min_frames": int(cal.get("codec_chunk_min_frames", 2)), + "codec_chunk_safety_margin_ms": float(cal.get("codec_chunk_safety_margin_ms", 50.0)), + } + if cal.get("codec_chunk_adaptive"): + out["codec_chunk_adaptive"] = True + return out diff --git a/vllm_omni/edge/hardware_probe.py b/vllm_omni/edge/hardware_probe.py new file mode 100644 index 000000000..de1e2a4d5 --- /dev/null +++ b/vllm_omni/edge/hardware_probe.py @@ -0,0 +1,292 @@ +"""Hardware capability probe for edge deployments. + +Produces a :class:`HardwareProfile` from ``/proc``, ``/sys`` and (optionally) +torch, and maps it to a coarse *hardware class* used to select deploy overlays +(``vllm_omni/deploy/edge/hardware/.yaml``) and to derive engine +overrides (see :mod:`vllm_omni.edge.adapt`). + +The probe reads from an injectable filesystem root so tests can feed captured +snapshots of other devices, and ``VLLM_OMNI_HW_PROFILE=`` replaces the +probe entirely (used for emulating a device class on a development host). +""" + +from __future__ import annotations + +import json +import os +import platform +import re +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +ENV_OVERRIDE = "VLLM_OMNI_HW_PROFILE" + +HW_CLASS_JETSON = "jetson" +HW_CLASS_CUDA_DISCRETE = "cuda_discrete" +HW_CLASS_ARM64_CPU = "arm64_cpu" +HW_CLASS_X86_CPU = "x86_cpu" +HW_CLASS_NPU_PHONE = "npu_phone" +HW_CLASSES = (HW_CLASS_JETSON, HW_CLASS_CUDA_DISCRETE, HW_CLASS_ARM64_CPU, HW_CLASS_X86_CPU, HW_CLASS_NPU_PHONE) + +# Flags we care about, normalized across x86 (``flags``) and ARM (``Features``). +_X86_FLAGS = ("avx2", "avx512f", "avx512_bf16", "avx512_vnni", "amx_bf16", "amx_tile", "amx_int8", "fma", "f16c") +_ARM_FLAGS = ("asimd", "asimdhp", "asimddp", "i8mm", "bf16", "sve", "sve2", "sme") +_FLAG_ALIASES = {"avx512f": "avx512", "asimd": "neon", "asimdhp": "fp16_arith", "asimddp": "dotprod"} + + +@dataclass +class HardwareProfile: + arch: str = "unknown" # x86_64 | aarch64 | ... + cpu_model: str = "" + cpu_flags: list[str] = field(default_factory=list) + n_logical: int = 0 + n_physical: int = 0 + clusters: list[dict[str, Any]] = field(default_factory=list) # [{"max_khz": int, "cpus": [int, ...]}], big first + numa_nodes: int = 1 + ram_total_bytes: int = 0 + ram_available_bytes: int = 0 + accelerator: str = "none" # none | cuda_discrete | cuda_unified | npu: + gpu_name: str = "" + gpu_mem_bytes: int = 0 + gpu_sm_count: int = 0 + gpu_capability: list[int] = field(default_factory=list) + supports_bf16: bool = False + supports_fp16: bool = True + thermal_zones: int = 0 + source: str = "probe" + + # ---------------------------------------------------------------- helpers + def has(self, *flags: str) -> bool: + return all(f in self.cpu_flags for f in flags) + + @property + def big_cpus(self) -> list[int]: + return list(self.clusters[0]["cpus"]) if self.clusters else list(range(self.n_logical)) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def to_json(self) -> str: + return json.dumps(self.to_dict(), indent=1) + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> HardwareProfile: + known = {f for f in cls.__dataclass_fields__} + return cls(**{k: v for k, v in d.items() if k in known}) + + @classmethod + def from_json(cls, text: str) -> HardwareProfile: + return cls.from_dict(json.loads(text)) + + +# ------------------------------------------------------------------ parsing +def _read(path: Path) -> str | None: + try: + return path.read_text(encoding="utf-8", errors="replace") + except OSError: + return None + + +def parse_cpuinfo(text: str) -> tuple[str, list[str], int]: + """Return (model name, normalized flags, logical cpu count) from /proc/cpuinfo.""" + model = "" + flags: set[str] = set() + n_logical = 0 + for line in text.splitlines(): + if ":" not in line: + continue + key, _, value = line.partition(":") + key = key.strip().lower() + value = value.strip() + if key == "processor": + n_logical += 1 + elif key in ("model name", "hardware", "cpu model") and not model: + model = value + elif key in ("flags", "features"): + for tok in value.split(): + if tok in _X86_FLAGS or tok in _ARM_FLAGS: + flags.add(_FLAG_ALIASES.get(tok, tok)) + return model, sorted(flags), n_logical + + +def parse_meminfo(text: str) -> tuple[int, int]: + total = avail = 0 + for line in text.splitlines(): + m = re.match(r"^(MemTotal|MemAvailable):\s+(\d+)\s*kB", line) + if m: + val = int(m.group(2)) * 1024 + if m.group(1) == "MemTotal": + total = val + else: + avail = val + return total, avail + + +def parse_clusters(cpu_max_khz: dict[int, int]) -> list[dict[str, Any]]: + """Group CPUs by max frequency (descending) -> big.LITTLE clusters.""" + groups: dict[int, list[int]] = {} + for cpu, khz in cpu_max_khz.items(): + groups.setdefault(khz, []).append(cpu) + return [{"max_khz": khz, "cpus": sorted(cpus)} for khz, cpus in sorted(groups.items(), key=lambda kv: -kv[0])] + + +def _sysfs_cpu_dirs(root: Path) -> list[Path]: + base = root / "sys" / "devices" / "system" / "cpu" + return sorted((p for p in base.glob("cpu[0-9]*") if p.is_dir()), key=lambda p: int(p.name[3:])) + + +def _physical_cores(root: Path, n_logical: int) -> int: + ids: set[tuple[str, str]] = set() + for cpu in _sysfs_cpu_dirs(root): + core = _read(cpu / "topology" / "core_id") + pkg = _read(cpu / "topology" / "physical_package_id") + if core is not None: + ids.add(((pkg or "0").strip(), core.strip())) + return len(ids) if ids else n_logical + + +def _cpu_max_khz(root: Path) -> dict[int, int]: + out: dict[int, int] = {} + for cpu in _sysfs_cpu_dirs(root): + raw = _read(cpu / "cpufreq" / "cpuinfo_max_freq") + if raw and raw.strip().isdigit(): + out[int(cpu.name[3:])] = int(raw.strip()) + return out + + +def _numa_nodes(root: Path) -> int: + base = root / "sys" / "devices" / "system" / "node" + nodes = [p for p in base.glob("node[0-9]*")] if base.exists() else [] + return max(1, len(nodes)) + + +def _thermal_zones(root: Path) -> int: + base = root / "sys" / "class" / "thermal" + return len(list(base.glob("thermal_zone*"))) if base.exists() else 0 + + +def _is_tegra(root: Path) -> bool: + return ( + (root / "etc" / "nv_tegra_release").exists() + or (root / "proc" / "device-tree" / "compatible").exists() + and "tegra" in (_read(root / "proc" / "device-tree" / "compatible") or "") + ) + + +def _torch_gpu_info() -> dict[str, Any]: + info: dict[str, Any] = {"available": False} + try: + import torch + + if not torch.cuda.is_available(): + return info + props = torch.cuda.get_device_properties(0) + info.update( + available=True, + name=props.name, + mem_bytes=int(props.total_memory), + sm_count=int(getattr(props, "multi_processor_count", 0)), + capability=[int(props.major), int(props.minor)], + bf16=bool(torch.cuda.is_bf16_supported()), + ) + except Exception: + return info + return info + + +# -------------------------------------------------------------------- probe +def probe(root: str | Path = "/", use_torch: bool = True) -> HardwareProfile: + """Probe the current machine (or a snapshot rooted at ``root``).""" + root = Path(root) + model, flags, n_logical = parse_cpuinfo(_read(root / "proc" / "cpuinfo") or "") + ram_total, ram_avail = parse_meminfo(_read(root / "proc" / "meminfo") or "") + arch = platform.machine() if root == Path("/") else _arch_from_flags(flags, model) + profile = HardwareProfile( + arch=arch, + cpu_model=model, + cpu_flags=flags, + n_logical=n_logical or os.cpu_count() or 1, + n_physical=_physical_cores(root, n_logical or os.cpu_count() or 1), + clusters=parse_clusters(_cpu_max_khz(root)), + numa_nodes=_numa_nodes(root), + ram_total_bytes=ram_total, + ram_available_bytes=ram_avail, + thermal_zones=_thermal_zones(root), + ) + gpu = _torch_gpu_info() if (use_torch and root == Path("/")) else {"available": False} + tegra = _is_tegra(root) + if gpu.get("available"): + profile.accelerator = "cuda_unified" if tegra else "cuda_discrete" + profile.gpu_name = gpu["name"] + profile.gpu_mem_bytes = gpu["mem_bytes"] + profile.gpu_sm_count = gpu["sm_count"] + profile.gpu_capability = gpu["capability"] + profile.supports_bf16 = gpu["bf16"] + elif tegra: + profile.accelerator = "cuda_unified" + profile.gpu_mem_bytes = ram_total # unified memory + profile.supports_bf16 = True + else: + profile.supports_bf16 = _cpu_bf16(profile) + profile.supports_fp16 = ( + True if profile.accelerator.startswith("cuda") else profile.has("fp16_arith") or arch == "x86_64" + ) + return profile + + +def _arch_from_flags(flags: list[str], model: str) -> str: + if any(f in flags for f in ("neon", "dotprod", "sve", "i8mm", "fp16_arith")): + return "aarch64" + if any(f in flags for f in ("avx2", "avx512", "amx_tile", "fma")): + return "x86_64" + return "unknown" + + +def _cpu_bf16(profile: HardwareProfile) -> bool: + if profile.arch == "x86_64": + return profile.has("avx512_bf16") or profile.has("amx_bf16") + if profile.arch == "aarch64": + return profile.has("bf16") + return False + + +def hardware_class(profile: HardwareProfile) -> str: + acc = profile.accelerator + if acc.startswith("npu"): + return HW_CLASS_NPU_PHONE + if acc == "cuda_unified": + return HW_CLASS_JETSON + if acc == "cuda_discrete": + return HW_CLASS_CUDA_DISCRETE + if profile.arch == "aarch64": + return HW_CLASS_ARM64_CPU + return HW_CLASS_X86_CPU + + +def load_profile(path: str | Path | None = None, use_torch: bool = True) -> HardwareProfile: + """Env override (``VLLM_OMNI_HW_PROFILE``) or explicit path wins over probing.""" + override = path or os.environ.get(ENV_OVERRIDE) + if override: + p = HardwareProfile.from_json(Path(override).read_text(encoding="utf-8")) + p.source = f"override:{override}" + return p + return probe(use_torch=use_torch) + + +def describe(profile: HardwareProfile) -> str: + cls = hardware_class(profile) + big = profile.big_cpus + return ( + f"class={cls} arch={profile.arch} cpu='{profile.cpu_model}' " + f"logical={profile.n_logical} physical={profile.n_physical} " + f"big_cpus={len(big)} flags={','.join(profile.cpu_flags)} ram={profile.ram_total_bytes / 2**30:.1f}GiB " + f"accel={profile.accelerator} gpu='{profile.gpu_name}' gpu_mem={profile.gpu_mem_bytes / 2**30:.1f}GiB " + f"bf16={profile.supports_bf16} source={profile.source}" + ) + + +if __name__ == "__main__": # pragma: no cover - manual use + prof = load_profile() + print(describe(prof)) + print(prof.to_json())