From 538eddd74d57ac5a8fea0a8ac0d5288e20678ccb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zhengyuan=20Su=20=28=E8=8B=8F=E6=94=BF=E6=B8=8A=29?= Date: Wed, 16 Sep 2026 20:49:04 +0800 Subject: [PATCH 01/99] [Perf][CosyVoice3] Bounded-window incremental HiFT vocoder streaming (#7521) Signed-off-by: Zhengyuan Su --- .../bench_cosyvoice3_concurrency_matrix.py | 517 ++++++++++++++++++ pyproject.toml | 2 - requirements/common.txt | 5 + .../cosyvoice3/test_cosyvoice3_components.py | 61 ++- .../test_cosyvoice3_incremental_hift.py | 415 ++++++++++++++ .../minicpmo_4_5/test_cuda_graph_wrapper.py | 2 + .../cosyvoice3/code2wav_core/hifigan.py | 170 +++++- .../models/cosyvoice3/cosyvoice3_code2wav.py | 136 +++-- 8 files changed, 1224 insertions(+), 84 deletions(-) create mode 100644 benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py diff --git a/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py new file mode 100644 index 000000000..fb45b1743 --- /dev/null +++ b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py @@ -0,0 +1,517 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix. + +Launches the server once per (concurrency, step-size) combo, streams TTS +requests cycling through several utterance lengths, and reports +TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is +``codec_chunk_frames``, the token-hop length the first emitted chunk waits +for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames`` +(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that +same ramp rather than from raw transport read boundaries, which don't align +to model chunks (see ``one_request``'s docstring). + +Running with no flags reproduces the default matrix. Per-request correctness +is checked: non-silent, long enough, and (with a fixed ``--seed``) +byte-identical PCM across requests for the same text. +""" + +import argparse +import json +import os +import signal +import statistics +import subprocess +import sys +import threading +import time +from pathlib import Path + +import httpx +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" + +DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav" +DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。" + +DEFAULT_TEXTS = [ + "Thank you so much for the birthday gift you sent me.", + "The weather is lovely today, with warm sunshine and a gentle breeze. " + "We decided to take a walk in the park to enjoy the spring flowers and the fresh air.", + "The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the " + "park to enjoy the spring flowers and the fresh air. Many people were out walking, children were " + "running and playing on the grass, and older folks sat on benches chatting with one another. We " + "strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting " + "on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we " + "continued on and reached a small lake, its water so clear you could see straight to the bottom, with " + "the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the " + "water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the " + "peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting " + "sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning " + "our next visit. On the way back we talked about everything we had seen and heard that day, feeling " + "wonderfully happy. It was truly an unforgettable weekend.", +] + +# PCM16 mono @ 24 kHz (CosyVoice3 output rate) +SAMPLE_RATE = 24000 + +# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte +# count and duration checks alone pass through. +MIN_RMS = 50.0 + +# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config) +DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit. +base_config: cosyvoice3.yaml + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: {chunk_frames} + codec_max_chunk_frames: {max_chunk_frames} + codec_stream_scale_factor: {stream_scale_factor} + codec_left_context_frames: 25 + codec_vocab_size: 6561 +""" + +# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens +# are consumed at token_mel_ratio mel frames each, and HiFT emits +# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the +# server's defaults) since one_request's byte-boundary reconstruction below +# depends on knowing them exactly. +TOKEN_MEL_RATIO = 2 +SAMPLES_PER_MEL = 480 +STREAM_SCALE_FACTOR = 2 + + +def _stats(values: list[float]) -> dict[str, float]: + """Mean/median/stdev/p95/min/max/n for a list of per-request measurements.""" + n = len(values) + if n == 0: + return { + "n": 0, + "mean": float("nan"), + "median": float("nan"), + "stdev": float("nan"), + "p95": float("nan"), + "min": float("nan"), + "max": float("nan"), + } + return { + "n": n, + "mean": statistics.mean(values), + "median": statistics.median(values), + "stdev": statistics.stdev(values) if n >= 2 else 0.0, + "p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0], + "min": min(values), + "max": max(values), + } + + +def vllm_bin() -> str: + for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"): + if c.is_file(): + return str(c) + return "vllm" + + +def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else [] + raise RuntimeError("server exited during startup:\n " + "\n ".join(tail)) + try: + if httpx.get(f"{url}/health", timeout=2.0).status_code == 200: + return + except Exception: + pass + time.sleep(2.0) + raise TimeoutError(f"server not ready within {timeout}s (log: {log})") + + +def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int): + """The token-hop-length ramp talker2code2wav_async_chunk uses to decide + when to fire the next flow-solve/emission: chunk_frames, then doubling + (STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not + codec_chunk_frames alone -- is what actually paces emitted audio chunks: + with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ... + + The very first emission is consistently one token short of chunk_frames + (confirmed against real speech.audio.delta event sizes) even though the + ramp's own progression is computed from the nominal chunk_frames, not the + actual first hop -- so only the first yield gets the -1 adjustment. + """ + yield chunk_frames - 1 + hop = chunk_frames + while True: + hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR)) + yield hop + + +def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None): + """Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes. + + ``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no + framing, so a client read boundary alone is a transport artifact, not a + model chunk boundary. Reconstructed instead by replicating the server's + own token-hop-length ramp (see ``_token_hop_lengths``) to compute each + real chunk's expected byte size ahead of time, and slicing the + accumulated byte stream at those offsets -- independent of how the + transport happens to fragment the underlying reads. The final chunk + (finalize, which releases extra held-back frames) is whatever remains + when the stream closes, not a predicted hop. + """ + payload = { + "model": model, + "input": text, + "response_format": "pcm", + "stream": True, + "stream_format": "audio", + "ref_audio": ref_audio, + "ref_text": ref_text, + } + if seed is not None: + payload["seed"] = seed + t0 = time.perf_counter() + ttfa = None + nbytes = 0 + audio_bytes = bytearray() + chunk_gaps_ms = [] # inter-chunk arrival gaps (per model chunk), by step index + prev_t = None + step_idx = 0 + + bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2 + hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames) + next_boundary = next(hop_gen) * bytes_per_token + last_crossed_bytes = 0 + received = 0 + last_now = t0 + + with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r: + if r.status_code != 200: + raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}") + for chunk in r.iter_bytes(): + if not chunk: + continue + now = time.perf_counter() + if ttfa is None: + ttfa = now - t0 + last_now = now + received += len(chunk) + nbytes += len(chunk) + audio_bytes += chunk + while received >= next_boundary: + if prev_t is not None: + chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3)) + step_idx += 1 + prev_t = now + last_crossed_bytes = next_boundary + next_boundary += next(hop_gen) * bytes_per_token + + if received > last_crossed_bytes: + # Finalize chunk: releases extra held-back frames, so it doesn't + # match a predicted hop size -- its arrival is the stream closing. + if prev_t is not None: + chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3)) + step_idx += 1 + + e2e = time.perf_counter() - t0 + audio_s = nbytes / 2 / SAMPLE_RATE + return { + "ttfa_s": ttfa, + "e2e_s": e2e, + "audio_s": audio_s, + "rtf": e2e / audio_s if audio_s > 0 else float("nan"), + "chunk_gaps_ms": chunk_gaps_ms, + "nbytes": nbytes, + "audio_bytes": audio_bytes, + } + + +def check_correctness(r: dict, min_audio_s: float) -> list[str]: + """Return a list of correctness violations (empty = pass).""" + problems = [] + if r["nbytes"] <= 0: + problems.append("empty audio (0 bytes)") + if r["audio_s"] < min_audio_s: + problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)") + if r["nbytes"] > 0: + samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64) + rms = float(np.sqrt(np.mean(samples**2))) + if rms < MIN_RMS: + problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})") + return problems + + +def drive(url, args, n, conc, label, results, lock, idx, chunk_frames): + """Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``.""" + texts = args.texts + + def worker(): + with httpx.Client(timeout=600.0) as c: + while True: + with lock: + i = idx[0] + idx[0] += 1 + if i >= n: + return + text = texts[i % len(texts)] + try: + r = one_request( + c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed + ) + r["text"] = text + with lock: + results.append(r) + except Exception as e: + print(f" [{label}] request {i} failed: {e}", flush=True) + + ts = [threading.Thread(target=worker) for _ in range(conc)] + t0 = time.perf_counter() + for t in ts: + t.start() + for t in ts: + t.join() + return time.perf_counter() - t0 + + +def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict: + tag = f"c{conc}_chunk{chunk_frames}" + run_dir = out / tag + run_dir.mkdir(parents=True, exist_ok=True) + log = run_dir / "server.log" + + deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml" + deploy.write_text( + DEPLOY_TEMPLATE.format( + chunk_frames=chunk_frames, + max_chunk_frames=4 * chunk_frames, + stream_scale_factor=STREAM_SCALE_FACTOR, + ) + ) + + cmd = [ + vllm_bin(), + "serve", + args.model, + "--omni", + "--deploy-config", + str(deploy), + "--host", + "127.0.0.1", + "--port", + str(args.port), + "--allowed-local-media-path", + str(REPO), + ] + env = dict(os.environ) + if args.cuda_visible_devices: + env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices + if args.trt is not None: + env["COSYVOICE3_TRT"] = args.trt + + print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True) + print(f" launching server (log: {log})", flush=True) + with open(log, "w") as lf: + proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True) + url = f"http://127.0.0.1:{args.port}" + try: + wait_ready(url, proc, log, args.server_timeout) + print(" server ready", flush=True) + + results: list[dict] = [] + lock = threading.Lock() + idx = [0] + if args.warmup: + print(f" warmup ({args.warmup} req)", flush=True) + drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames) + results.clear() + idx[0] = 0 + + print(f" running {args.num_prompts} req @ concurrency {conc}", flush=True) + wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames) + + ok = [r for r in results if r.get("ttfa_s") is not None] + problems = [] + for r in ok: + problems += check_correctness(r, args.min_audio_s) + # With a fixed seed, same-text requests must produce byte-identical PCM (per text group) + if args.seed is not None: + by_text: dict[str, bytes] = {} + for r in ok: + t = r.get("text") + ref = by_text.get(t) + if ref is None: + by_text[t] = r.get("audio_bytes") + elif r.get("audio_bytes") != ref: + problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM") + n_ok = len(ok) + total_audio_s = sum(r["audio_s"] for r in ok) + all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]] # (step_index, gap_ms) per request + gap_vals = [g for _, g in all_gaps] + + # bucket gaps by chunk position within a request, to see any O(n) slope + by_idx: dict[int, list[float]] = {} + for step_index, g in all_gaps: + by_idx.setdefault(step_index, []).append(g) + step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())} + + agg = { + "concurrency": conc, + "step_size": chunk_frames, + "ok": n_ok, + "errors": args.num_prompts - n_ok, + "correctness_problems": len(problems), + "correctness_sample": problems[:5], + "total_audio_s": total_audio_s, + "throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"), + "ttfa_s": _stats([r["ttfa_s"] for r in ok]), + "e2e_s": _stats([r["e2e_s"] for r in ok]), + "rtf": _stats([r["rtf"] for r in ok]), + "step_latency_ms": _stats(gap_vals), + "step_latency_ms_by_idx": step_latency_by_idx, + "num_steps": len(all_gaps), + # per-utterance-length breakdown; n per text = num_prompts / len(texts) + "by_text": { + t: { + "audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]), + "e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]), + "rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]), + } + for t in sorted({r.get("text") for r in ok}) + }, + } + print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True) + return agg + finally: + print(" shutting down server", flush=True) + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except Exception: + proc.terminate() + try: + proc.wait(timeout=60) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except Exception: + proc.kill() + proc.wait(timeout=20) + deploy.unlink(missing_ok=True) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", default=DEFAULT_MODEL) + p.add_argument( + "--text", + dest="texts", + nargs="+", + default=DEFAULT_TEXTS, + help="One or more utterance texts. With several, requests cycle through " + "them so a single server launch covers all utterance lengths; results " + "are bucketed per text.", + ) + p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO) + p.add_argument("--ref-text", default=DEFAULT_REF_TEXT) + p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8]) + p.add_argument( + "--step-size", + type=int, + nargs="+", + default=[15], + help="codec_chunk_frames values to sweep (the streaming step size).", + ) + p.add_argument("--num-prompts", type=int, default=30) + p.add_argument("--warmup", type=int, default=2) + p.add_argument("--port", type=int, default=8091) + p.add_argument( + "--server-timeout", + type=float, + default=1800.0, + help="Startup budget; first run also downloads weights and builds the TRT engine.", + ) + p.add_argument("--output-dir", default="./perf_steps") + p.add_argument("--cuda-visible-devices", default="0") + p.add_argument( + "--trt", + choices=["0", "1"], + default=None, + help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.", + ) + p.add_argument( + "--min-audio-s", + type=float, + default=1.0, + help="Minimum expected audio seconds per request for the correctness check.", + ) + p.add_argument( + "--seed", + type=int, + default=42, + help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' " + "equivalent (any non-int) to disable -- there is no unset sentinel here " + "by design, since the default sweep is meant to be exactly reproducible.", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + out = Path(args.output_dir).resolve() + out.mkdir(parents=True, exist_ok=True) + + all_results = [] + failed = False + for conc in args.concurrency: + for chunk_frames in args.step_size: + try: + result = run_combo(args, conc, chunk_frames, out) + all_results.append(result) + if result["errors"] > 0 or result["correctness_problems"] > 0: + failed = True + except Exception as e: + print(f" combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True) + failed = True + + out_path = out / "steps_results.json" + with open(out_path, "w", encoding="utf-8") as f: + json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False) + print(f"\nResults written to {out_path}", flush=True) + + print("\n=== SUMMARY ===") + header = ( + f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} " + f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}" + ) + print(header) + print("-" * len(header)) + for r in all_results: + print( + f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} " + f"{r['ttfa_s']['median'] * 1e3:>12.1f} " + f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} " + f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}" + ) + for t, stats in r.get("by_text", {}).items(): + print( + f" text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}" + f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms (n={stats['e2e_s']['n']}, " + f"audio {stats['audio_s']['mean']:.1f}s)" + ) + + if failed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 434c4305c..c463d3917 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,14 +66,12 @@ dev = [ "jiwer==4.0.0", "zhon==2.1.1", "zhconv==1.4.3", - "scipy==1.18.0", "funasr==1.3.29", "lpips==0.1.4", # for quantization fp8 "orjson==3.11.9", # to make v1/embeddings API fast "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests "torchdiffeq==0.2.5", # for covo_audio Test "FlagEmbedding==1.4.0", # for hunyuan_image3 tests - "s3tokenizer==0.3.0", # for tts test # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, # whose top-level stepaudio2 package conflicts with this distribution. "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) diff --git a/requirements/common.txt b/requirements/common.txt index 855fb6b0a..d7b845e32 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -41,3 +41,8 @@ pydantic-settings>=2.0.0 gguf>=0.10.0 # Required for running the Cosmos3 model with guardrails enabled (default behavior). cosmos-guardrail>=0.3.1 +# Required for CosyVoice3 (get_window/resample_poly). >=1.18.0 requires +# Python >=3.12, breaking base installs on the project's supported 3.10/3.11; +# these are long-stable APIs with no version-specific need for 1.18. +scipy>=1.11.0 +s3tokenizer==0.3.0 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 7dfbe6e74..74fc5059b 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -15,6 +15,7 @@ from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( CausalHiFTGenerator, + HiFTGenerator, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -76,6 +77,39 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift): assert causal_hift.stft_window.device == waveform.device +class _MinimalF0Predictor(nn.Module): + def __init__(self) -> None: + super().__init__() + self.proj = nn.Conv1d(80, 1, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x).squeeze(1).abs() + + +@pytest.mark.parametrize("sampling_rate", [22050, 24000]) +def test_noncausal_hift_inference_runs_with_real_sinegen(sampling_rate): + """Base (non-causal) HiFTGenerator with a real SineGen/SineGen2 (not a + test double) must still run: SourceModuleHnNSF/SineGen/SineGen2/_f02sine + all branch on `self.causal` to return the pre-streaming (shorter) tuple + contract for non-causal callers, and any one of them regressing back to + always returning the streaming-length tuple breaks this immediately. + """ + hift = HiFTGenerator( + base_channels=32, + sampling_rate=sampling_rate, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=_MinimalF0Predictor(), + ).eval() + + speech, _ = hift.inference(torch.randn(1, 80, 20)) + + assert speech.shape[0] == 1 + assert torch.isfinite(speech).all() + + def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch): class NpuTensor: device = SimpleNamespace(type="npu") @@ -83,16 +117,17 @@ def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypa tensor = NpuTensor() stft_calls = [] istft_calls = [] - monkeypatch.setattr( - causal_hift, - "_stft_on_cpu", - lambda value: stft_calls.append(value) or ("real", "imag"), - ) - monkeypatch.setattr( - causal_hift, - "_istft_on_cpu", - lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform", - ) + + def fake_stft(value): + stft_calls.append(value) + return ("real", "imag") + + def fake_istft(magnitude, phase): + istft_calls.append((magnitude, phase)) + return "waveform" + + monkeypatch.setattr(causal_hift, "_stft_on_cpu", fake_stft) + monkeypatch.setattr(causal_hift, "_istft_on_cpu", fake_istft) assert causal_hift._stft(tensor) == ("real", "imag") assert causal_hift._istft(tensor, tensor) == "waveform" @@ -627,7 +662,11 @@ def test_code2wav_forward_finalizes_hift_tail(): def inference(self, speech_feat, finalize=True): self.finalize_calls.append(bool(finalize)) - return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None + return ( + torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), + None, + None, + ) model = object.__new__(CosyVoice3Code2Wav) nn.Module.__init__(model) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py new file mode 100644 index 000000000..1e8eff599 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Correctness + boundedness tests for the incremental (windowed) HiFT path. + +The streaming vocoder runs HiFT over a bounded mel window instead of the full +cumulative history, carrying harmonic phase and noise-buffer offset across +chunks. + +* **Correctness** — windowed output must match the full cumulative re-run. + With full history the match is bit-exact (proves the phase/noise carry is + right). With a bounded window, float64 still agrees to ~1e-16 (the window + is algebraically exact); float32 differs by ~2e-7 absolute (reassociation + rounding, not truncation), so tests assert a tight atol/rtol rather than + PCM16 byte-identity, since a sub-LSB perturbation can flip a boundary sample. +* **Boundedness** — the per-chunk window must not grow with mel history. + +Tests are parametrized across REAL_CONFIG (24000 Hz -> SineGen2, what +CosyVoice3 ships) and LEGACY_CONFIG (22050 Hz -> SineGen) so a class-selection +regression in either path is caught. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + CausalConvRNNF0Predictor, + CausalHiFTGenerator, + SineGen, + SineGen2, + _wrapped_slice, +) +from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +CHUNK_LEN = 24 +TOTAL_MEL = 96 +SPM = 480 # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model + +LONG_TOTAL_MEL = 400 # long enough that the 64-frame window actually truncates +PCM16_LSB = 1.0 / 32767.0 # smallest representable PCM16 difference + +# Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom. +ATOL = 1e-6 +RTOL = 1e-5 + +# What CosyVoice3 ships (-> SineGen2). nb_harmonics=8 is required, not a free +# choice: SineGen2's causal buffers are hardcoded to dim 9 (= harmonic_num+1). +REAL_CONFIG = {"sampling_rate": 24000, "nb_harmonics": 8} +LEGACY_CONFIG = {"sampling_rate": 22050, "nb_harmonics": 4} # -> SineGen, still shipped code + +CONFIGS = [ + pytest.param(REAL_CONFIG, id="sinegen2_real_config"), + pytest.param(LEGACY_CONFIG, id="sinegen1_legacy_config"), +] + + +def _make_hift(config: dict = REAL_CONFIG) -> CausalHiFTGenerator: + """Small CausalHiFTGenerator; defaults to the config CosyVoice3 ships.""" + torch.manual_seed(0) + f0_predictor = CausalConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=16) + return CausalHiFTGenerator( + in_channels=80, + base_channels=32, + nb_harmonics=config["nb_harmonics"], + sampling_rate=config["sampling_rate"], + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + resblock_kernel_sizes=[3, 7, 11], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + f0_predictor=f0_predictor, + ).eval() + + +def test_real_config_resolves_to_sinegen2(): + """Guard against silently testing the wrong SineGen class again.""" + hift = _make_hift(REAL_CONFIG) + assert isinstance(hift.m_source.l_sin_gen, SineGen2) + + hift_legacy = _make_hift(LEGACY_CONFIG) + assert isinstance(hift_legacy.m_source.l_sin_gen, SineGen) + + +def _make_model(hift: CausalHiFTGenerator, window_len: int) -> CosyVoice3Code2Wav: + """Build a CosyVoice3Code2Wav shell wired to the given HiFT.""" + model = object.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.hift = hift + model._hift_window_len = window_len + return model + + +def _chunks(total_mel: int = TOTAL_MEL, chunk_len: int = CHUNK_LEN) -> list[torch.Tensor]: + torch.manual_seed(1) + mel = torch.randn(1, 80, total_mel) + return [mel[:, :, i : i + chunk_len] for i in range(0, total_mel, chunk_len)] + + +def _variable_chunks(total_mel: int, first_len: int, step_len: int) -> list[torch.Tensor]: + """First chunk sized like a real prompt prefill, then small fixed-size steps.""" + torch.manual_seed(1) + mel = torch.randn(1, 80, total_mel) + chunks = [mel[:, :, :first_len]] + pos = first_len + while pos < total_mel: + chunks.append(mel[:, :, pos : pos + step_len]) + pos += step_len + return chunks + + +def _full_reference(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor: + """Re-run HiFT over the full cumulative mel each chunk (pre-windowing behavior). + + Uses non-finalize inference per chunk, matching the old streaming vocoder; a + single finalize pass would emit the look-right tail streaming holds back. + """ + torch.manual_seed(0) + emitted = [] + cache: dict[str, torch.Tensor] | None = None + for chunk in chunks: + cached = None if cache is None else cache.get("mel") + if cached is not None and cached.numel() > 0: + tts_mel = torch.cat([cached, chunk], dim=-1) + else: + tts_mel = chunk + speech, _, _ = model.hift.inference(speech_feat=tts_mel, finalize=False) + speech = speech.reshape(speech.shape[0], -1) + offset = 0 if cache is None else cache.get("speech_offset") + emitted.append(speech[:, offset:]) + cache = {"mel": tts_mel.detach().cpu(), "speech_offset": int(speech.shape[-1])} + return torch.cat(emitted, dim=-1) + + +def _pcm16(x: torch.Tensor) -> torch.Tensor: + """Quantize float audio to PCM16, the format the server streams to clients.""" + return (x.clamp(-1, 1) * 32767).round().to(torch.int16) + + +def _incremental(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor: + """Run the new windowed path via _stream_hift_from_feat.""" + torch.manual_seed(0) + emitted = [] + cache: dict[str, torch.Tensor] | None = None + for chunk in chunks: + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + emitted.append(speech.reshape(speech.shape[0], -1)) + return torch.cat(emitted, dim=-1) + + +def _rel_divergence(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a - b).abs().mean() / a.abs().mean().clamp_min(1e-6)).item() + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_full_reference_exactly_with_full_history(config): + """With a window covering the full history, the windowed path must match the + full cumulative re-run exactly. This proves the phase/noise carry is right.""" + hift = _make_hift(config) + model = _make_model(hift, window_len=TOTAL_MEL) # window >= history + chunks = _chunks() + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}" + assert _rel_divergence(full, incr) < 1e-6 + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_bounded_window_is_close(config): + """With a bounded window, the windowed path stays close to the full re-run; + only the truncated deep history diverges, and it is small.""" + hift = _make_hift(config) + model = _make_model(hift, window_len=48) + chunks = _chunks() + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_streaming_reference_within_tolerance(config): + """Windowed output matches the full-cumulative reference within float32 rounding. + + The load-bearing correctness test: 400 mel frames far exceeds the 64-frame + window, so truncation genuinely exercises the phase/noise carry. Float64 + agrees to ~1e-16, confirming the residual here is rounding, not truncation. + """ + hift = _make_hift(config) + model = _make_model(hift, window_len=64) # the real _hift_window_len + chunks = _chunks(total_mel=LONG_TOTAL_MEL) + + trim = int(hift.f0_predictor.condnet[0].causal_padding) + assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN # window must actually truncate + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}" + torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL) + + max_dev = (full - incr).abs().max().item() # stays far below one PCM16 step + assert max_dev < PCM16_LSB / 10, f"max deviation {max_dev:.3e} vs PCM16 LSB {PCM16_LSB:.3e}" + + +def test_incremental_hift_window_is_bounded(): + """Per-chunk window must not grow with cumulative mel history.""" + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=200, chunk_len=24) + + trim = int(hift.f0_predictor.condnet[0].causal_padding) + f0_margin = int(hift.f0_predictor.left_context_frames) + + cache = None + window_sizes = [] + for chunk in chunks: + _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + window_sizes.append(int(cache["mel"].shape[-1])) + + assert max(window_sizes) <= 32 + trim + f0_margin + 24 # bounded regardless of utterance length + steady = max(window_sizes) + assert window_sizes.count(steady) >= 2 # steady-state value repeats once history exceeds window + + +def test_incremental_hift_finalize_releases_tail(): + """Finalize must emit the released look-right tail without crashing.""" + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=72, chunk_len=24) + + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + if finalize: + assert cache is None + + full = torch.cat(emitted, dim=-1) + assert full.shape[-1] > 0 + assert torch.isfinite(full).all() + + +def test_incremental_hift_emits_full_audio_length(): + """Streamed output covers the full mel history at samples-per-mel resolution. + + Guards against emission arithmetic dropping an upsample stage, which + silently truncates audio (e.g. 1/3 of expected length with [8,5,3]). + """ + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=72, chunk_len=24) + + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + if finalize: + assert cache is None + + full = torch.cat(emitted, dim=-1) + # total_mel * SPM, +slack for finalize's look-right tail, -slack for a dropped upsample stage + total_mel = sum(c.shape[-1] for c in chunks) + assert full.shape[-1] > total_mel * SPM * 0.8, f"{full.shape[-1]} vs {total_mel * SPM}" + assert full.shape[-1] < total_mel * SPM * 1.5, f"{full.shape[-1]} vs {total_mel * SPM}" + + +@pytest.mark.parametrize("config", CONFIGS) +def test_finalize_matches_single_pass_length_exactly(config): + """Streamed total (chunked, last=finalize) must exactly match a single + non-streaming finalize=True call over the whole mel -- not just approximately. + + The finalize release must include `trim` (the F0 predictor's own held-back + frames), or the last 3 mel frames (1,440 samples / 60ms at 24kHz) of every + stream's audio go missing. + """ + hift = _make_hift(config) + chunks = _chunks(total_mel=200, chunk_len=24) + + torch.manual_seed(0) + full_mel = torch.cat(chunks, dim=-1) + ground_truth, _, _ = hift.inference(speech_feat=full_mel, finalize=True) + ground_truth = ground_truth.reshape(ground_truth.shape[0], -1) + + hift2 = _make_hift(config) + model = _make_model(hift2, window_len=64) + torch.manual_seed(0) + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + streamed = torch.cat(emitted, dim=-1) + + assert streamed.shape[-1] == ground_truth.shape[-1], ( + f"{streamed.shape[-1]} vs {ground_truth.shape[-1]} (diff={ground_truth.shape[-1] - streamed.shape[-1]} samples)" + ) + torch.testing.assert_close(streamed, ground_truth, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_reference_for_voiced_f0(config): + """Forces F0 above nsf_voiced_threshold (untrained weights never trigger voiced, + so no other test here does) and checks windowed output against the full- + cumulative reference, which exercises condnet's left-causal convs and + SineGen/SineGen2's phase integration under genuinely voiced harmonics. + + Tolerance is looser than ATOL/RTOL: summing this sustained constant-F0 phase + in per-window chunks vs. one recompute reassociates float32 adds differently; + real speech's varying/unvoiced F0 keeps this far smaller in practice. + """ + hift = _make_hift(config) + with torch.no_grad(): + hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10) + chunks = _chunks(total_mel=200, chunk_len=24) + + full = _full_reference(_make_model(hift, window_len=400), chunks) + incr = _incremental(_make_model(hift, window_len=64), chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4) + + +@pytest.mark.parametrize("config", CONFIGS) +@pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)]) +def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len): + """Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small, + sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform + 24-frame chunks used elsewhere. cache_state must carry the F0 predictor's own + left receptive field across calls regardless of how small a single chunk is, + or the next call's margin comes up short and substitutes zero-padding for real + history -- corrupting SineGen's phase integration the same way as an absent + margin entirely. + """ + hift = _make_hift(config) + with torch.no_grad(): + hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10) + chunks = _variable_chunks(total_mel=200, first_len=first_len, step_len=step_len) + + full = _full_reference(_make_model(hift, window_len=400), chunks) + incr = _incremental(_make_model(hift, window_len=64), chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4) + + +def test_wrapped_slice_wraps_position_deterministically(): + """SineGen/SineGen2/SourceModuleHnNSF's noise buffers are a fixed 300s at + 24kHz; a stream running past that must keep repeating the realization + (position-deterministic) rather than the slice coming back short and + failing to broadcast against the rest of that call's tensors. + """ + buf = torch.arange(10).reshape(1, 10, 1) + + assert _wrapped_slice(buf, 2, 3).flatten().tolist() == [2, 3, 4] + assert _wrapped_slice(buf, 8, 4).flatten().tolist() == [8, 9, 0, 1] + assert _wrapped_slice(buf, 12, 3).flatten().tolist() == [2, 3, 4] # offset past the end + assert _wrapped_slice(buf, 5, 13).flatten().tolist() == [5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7] + + +def test_incremental_hift_empty_middle_chunk_emits_nothing(): + """An empty middle chunk (reachable from forward_streaming_batch when a + row's post-trim feat comes out empty near the lookahead boundary) must + emit zero samples and leave cache_state's phase_acc unchanged, not + re-emit the re-decoded overlap window as if it were new audio. + """ + hift = _make_hift(REAL_CONFIG) + model = _make_model(hift, window_len=64) + chunks = _chunks(total_mel=96, chunk_len=24) + empty = chunks[0][:, :, :0] + + cache = None + for chunk in chunks[:2]: + _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + assert cache is not None + phase_acc_before = cache["phase_acc"] + + speech, cache = model._stream_hift_from_feat(empty, cache_state=cache, finalize=False) + + assert speech.shape[-1] == 0 + if phase_acc_before is None: + assert cache["phase_acc"] is None + else: + torch.testing.assert_close(cache["phase_acc"], phase_acc_before) + + +@pytest.mark.parametrize("config", CONFIGS) +@pytest.mark.parametrize("mel_frames", [1, 2, 3, 4]) +def test_incremental_hift_short_finalize_chunk_does_not_crash(config, mel_frames): + """A one-token (or otherwise sub-trim) utterance that finalizes on its very + first chunk must not crash. On finalize the F0 predictor consumes the whole + window (it does not hold back `trim` frames), so the phase-carry index must + not be inflated by `trim` -- it would point past the end of a short window. + """ + hift = _make_hift(config) + model = _make_model(hift, window_len=32) + chunk = torch.randn(1, 80, mel_frames) + + speech, cache = model._stream_hift_from_feat(chunk, cache_state=None, finalize=True) + + assert cache is None + assert speech.shape[0] == 1 + assert speech.shape[-1] > 0 + assert torch.isfinite(speech).all() diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 959156b8f..555e72e46 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -177,6 +177,8 @@ class _MiniDiT(nn.Module): cnn_cache_buffer: torch.Tensor | None = None, att_cache_buffer: torch.Tensor | None = None, ) -> torch.Tensor: + assert cnn_cache_buffer is not None + assert att_cache_buffer is not None x = x.transpose(1, 2) x = self.in_proj(x) for b_idx in range(len(self.blocks)): diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 7cbd31e9c..2b8f3e61e 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -98,6 +98,46 @@ class ResBlock(torch.nn.Module): remove_weight_norm(self.convs2[idx]) +def _carry_phase_at_boundary( + cum_total: torch.Tensor, + phase_acc: torch.Tensor | None, + next_overlap: int, + *, + dim: int, + trim: int = 0, +) -> torch.Tensor | None: + """Accumulated phase (mod 1) carried into the next streaming window. + + Shared by SineGen and SineGen2 so harmonic phase stays continuous across + chunk boundaries when HiFT runs over a bounded mel window. ``cum_total`` + must already include ``phase_acc``; falls back to it unchanged if + ``next_overlap`` (+ ``trim``) covers the whole window. + """ + length = cum_total.shape[dim] + if next_overlap >= length + trim: + return phase_acc + idx = length + trim - next_overlap - 1 + return cum_total.narrow(dim, idx, 1) % 1 + + +def _wrapped_slice(buf: torch.Tensor, offset: int, length: int, dim: int = 1) -> torch.Tensor: + """``length`` elements of ``buf`` along ``dim`` starting at ``offset``, + wrapping (mod ``buf.shape[dim]``) instead of running past the end and + coming back short. Shared by SineGen/SineGen2/SourceModuleHnNSF's + position-indexed noise buffers, now that streams can outlast their 300s. + """ + n = buf.shape[dim] + start = offset % n + parts = [] + remaining = length + while remaining > 0: + take = min(remaining, n - start) + parts.append(buf.narrow(dim, start, take)) + remaining -= take + start = 0 + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim) + + class SineGen(torch.nn.Module): """Definition of sine generator SineGen(samp_rate, harmonic_num = 0, @@ -114,13 +154,24 @@ class SineGen(torch.nn.Module): segment is always sin(np.pi) or cos(0) """ - def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0): + def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0, causal=False): super().__init__() self.sine_amp = sine_amp self.noise_std = noise_std self.harmonic_num = harmonic_num self.sampling_rate = samp_rate self.voiced_threshold = voiced_threshold + self.causal = causal + # Fixed per model instance like SineGen2's rand_ini, not redrawn per + # call, or each harmonic jumps to a new random phase every window. + self.phase_vec = Uniform(low=-np.pi, high=np.pi).sample(sample_shape=(1, harmonic_num + 1, 1)) + self.phase_vec[:, 0, :] = 0 + if causal: + self.register_buffer( + "noise_buf", + torch.randn(1, 300 * 24000, harmonic_num + 1), + persistent=False, + ) def _f02uv(self, f0): # generate uv signal @@ -128,22 +179,32 @@ class SineGen(torch.nn.Module): return uv @torch.no_grad() - def forward(self, f0): + def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0): """sine_tensor, uv = forward(f0) input F0: tensor(batchsize=1, dim=1, length) f0 for unvoiced steps should be 0 output sine_tensor: tensor(batchsize=1, length, dim) output uv: tensor(batchsize=1, length, 1) + + ``next_overlap`` and ``trim`` are sample-indexed here (``f0`` arrives + post-``f0_upsamp``), unlike SineGen2 which downsamples to mel rate + before its own cumsum -- callers must scale both by samples-per-mel. """ f0 = f0.transpose(1, 2) F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) for i in range(self.harmonic_num + 1): F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate - theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) - u_dist = Uniform(low=-np.pi, high=np.pi) - phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) - phase_vec[:, 0, :] = 0 + cum = torch.cumsum(F_mat, dim=-1) + cum_total = cum if phase_acc is None else phase_acc + cum + theta_mat = 2 * np.pi * (cum_total % 1) + new_phase_acc = _carry_phase_at_boundary(cum_total, phase_acc, next_overlap, dim=-1, trim=trim) + if self.causal: + phase_vec = self.phase_vec.to(F_mat.device) + else: + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 # generate sine waveforms sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) @@ -155,11 +216,16 @@ class SineGen(torch.nn.Module): # std = self.sine_amp/3 -> max value ~ self.sine_amp # . for voiced regions is self.noise_std noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 - noise = noise_amp * torch.randn_like(sine_waves) + if self.training is False and self.causal is True: + noise = noise_amp * _wrapped_slice(self.noise_buf, noise_offset, sine_waves.shape[-1]).transpose(1, 2) + else: + noise = noise_amp * torch.randn_like(sine_waves) # first: set the unvoiced part to 0 by uv # then: additive noise sine_waves = sine_waves * uv + noise + if self.causal: + return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise, new_phase_acc return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise @@ -215,7 +281,7 @@ class SineGen2(torch.nn.Module): uv = (f0 > self.voiced_threshold).type(torch.float32) return uv - def _f02sine(self, f0_values): + def _f02sine(self, f0_values, phase_acc=None, next_overlap=0, trim=0): """f0_values: (batchsize, length, dim) where dim indicates fundamental tone and overtones """ @@ -223,9 +289,12 @@ class SineGen2(torch.nn.Module): # because 2 * np.pi * n doesn't affect phase rad_values = (f0_values / self.sampling_rate) % 1 - # initial phase noise (no noise for fundamental component) + # initial phase noise (no noise for fundamental component). Only applied + # on the genuine first chunk of a stream (see rationale above the class); + # phase_acc is None exactly when there is no carried-forward phase yet. if self.training is False and self.causal is True: - rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) + if phase_acc is None: + rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) else: rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) rand_ini[:, 0] = 0 @@ -237,7 +306,10 @@ class SineGen2(torch.nn.Module): rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear" ).transpose(1, 2) - phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + cum = torch.cumsum(rad_values, dim=1) + phase = cum if phase_acc is None else cum + phase_acc + new_phase_acc = _carry_phase_at_boundary(phase, phase_acc, next_overlap, dim=1, trim=trim) + phase = phase * 2 * np.pi phase = torch.nn.functional.interpolate( phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, @@ -272,9 +344,10 @@ class SineGen2(torch.nn.Module): # get the sines sines = torch.cos(i_phase * 2 * np.pi) - return sines + new_phase_acc = None + return (sines, new_phase_acc) if self.causal else sines - def forward(self, f0): + def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0): """sine_tensor, uv = forward(f0) input F0: tensor(batchsize=1, length, dim=1) f0 for unvoiced steps should be 0 @@ -284,8 +357,12 @@ class SineGen2(torch.nn.Module): # fundamental component fn = torch.multiply(f0, self.harmonic_ids) - # generate sine waveforms - sine_waves = self._f02sine(fn) * self.sine_amp + if self.causal: + sine_waves, new_phase_acc = self._f02sine(fn, phase_acc, next_overlap, trim) + else: + sine_waves = self._f02sine(fn) + new_phase_acc = None + sine_waves = sine_waves * self.sine_amp # generate uv signal uv = self._f02uv(f0) @@ -295,14 +372,15 @@ class SineGen2(torch.nn.Module): # . for voiced regions is self.noise_std noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 if self.training is False and self.causal is True: - noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device) + # indexed by absolute sample position, not window-relative position + noise = noise_amp * _wrapped_slice(self.sine_waves, noise_offset, sine_waves.shape[1]).to(sine_waves.device) else: noise = noise_amp * torch.randn_like(sine_waves) # first: set the unvoiced part to 0 by uv # then: additive noise sine_waves = sine_waves * uv + noise - return sine_waves, uv, noise + return (sine_waves, uv, noise, new_phase_acc) if self.causal else (sine_waves, uv, noise) class SourceModuleHnNSF(torch.nn.Module): @@ -338,10 +416,13 @@ class SourceModuleHnNSF(torch.nn.Module): self.sine_amp = sine_amp self.noise_std = add_noise_std + self.upsample_scale = upsample_scale # to produce sine waveforms if sinegen_type == "1": - self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold) + self.l_sin_gen = SineGen( + sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal + ) else: self.l_sin_gen = SineGen2( sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal @@ -354,24 +435,38 @@ class SourceModuleHnNSF(torch.nn.Module): if causal is True: self.uv = torch.rand(1, 300 * 24000, 1) - def forward(self, x): + def forward(self, x, phase_acc=None, uv_offset=0, next_overlap=0, trim=0): """ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) F0_sampled (batchsize, length, 1) Sine_source (batchsize, length, 1) noise_source (batchsize, length 1) """ - # source for harmonic branch with torch.no_grad(): - sine_wavs, uv, _ = self.l_sin_gen(x) + if not self.causal: + sine_wavs, uv, _ = self.l_sin_gen(x) + new_phase_acc = None + elif isinstance(self.l_sin_gen, SineGen): + # next_overlap/trim arrive in mel frames; SineGen needs samples. + sine_wavs, uv, _, new_phase_acc = self.l_sin_gen( + x, + phase_acc, + next_overlap * self.upsample_scale, + trim * self.upsample_scale, + noise_offset=uv_offset, + ) + else: + sine_wavs, uv, _, new_phase_acc = self.l_sin_gen( + x, phase_acc, next_overlap, trim, noise_offset=uv_offset + ) sine_merge = self.l_tanh(self.l_linear(sine_wavs)) # source for noise branch, in the same shape as uv if self.training is False and self.causal is True: - noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3 + noise = _wrapped_slice(self.uv, uv_offset, uv.shape[1]) * self.sine_amp / 3 else: noise = torch.randn_like(uv) * self.sine_amp / 3 - return sine_merge, noise, uv + return sine_merge, noise, uv, new_phase_acc class HiFTGenerator(nn.Module): @@ -606,7 +701,7 @@ class HiFTGenerator(nn.Module): f0 = self.f0_predictor(speech_feat) # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, _ = self.m_source(s) s = s.transpose(1, 2) # mel+source->speech generated_speech = self.decode(x=speech_feat, s=s) @@ -622,7 +717,7 @@ class HiFTGenerator(nn.Module): f0 = self.f0_predictor(speech_feat) # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, _ = self.m_source(s) s = s.transpose(1, 2) # use cache_source to avoid glitch if cache_source.shape[2] != 0: @@ -785,14 +880,28 @@ class CausalHiFTGenerator(HiFTGenerator): return x @torch.inference_mode() - def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor: + def inference( + self, + speech_feat: torch.Tensor, + finalize: bool = True, + phase_acc: torch.Tensor | None = None, + uv_offset: int = 0, + next_overlap: int = 0, + trim: int = 0, + f0_margin: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move # self.f0_predictor to cpu if necessary self.f0_predictor.to("cpu") f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat) + if f0_margin > 0: + f0 = f0[:, f0_margin:] + speech_feat = speech_feat[:, :, f0_margin:] # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, new_phase_acc = self.m_source( + s, phase_acc=phase_acc, uv_offset=uv_offset, next_overlap=next_overlap, trim=trim + ) s = s.transpose(1, 2) if finalize is True: generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize) @@ -800,7 +909,7 @@ class CausalHiFTGenerator(HiFTGenerator): generated_speech = self.decode( x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize ) - return generated_speech, s + return generated_speech, s, new_phase_acc class CausalConv1dUpsample(torch.nn.Conv1d): @@ -954,6 +1063,11 @@ class CausalConvRNNF0Predictor(nn.Module): nn.ELU(), ) self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + self.left_context_frames = sum( + layer.causal_padding + for layer in self.condnet + if isinstance(layer, CausalConv1d) and layer.causal_type == "left" + ) def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: if finalize is True: diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index ef706b6c8..1713467c8 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -9,10 +9,8 @@ This module contains the code2wav (token-to-waveform) stage which uses: 3. HiFiGAN vocoder for waveform synthesis """ -from __future__ import annotations - from collections import Counter -from typing import cast +from typing import TypedDict, cast import numpy as np import torch @@ -39,6 +37,21 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) +class StreamingFlowItem(TypedDict, total=False): + """One entry of the batched-streaming item list passed to forward_streaming_batch.""" + + index: int + req_id: str | None + stream_finished: bool + token: torch.Tensor + prompt_token: torch.Tensor + prompt_feat: torch.Tensor + embedding: torch.Tensor + cache_state: dict[str, torch.Tensor] | None + token_offset_tokens: int + finalize: bool + + class CosyVoice3Code2Wav(nn.Module): """CosyVoice3 Code2Wav stage for token-to-waveform conversion. @@ -119,6 +132,10 @@ class CosyVoice3Code2Wav(nn.Module): self.mel_cache_len = 20 self.source_cache_len = int(self.mel_cache_len * 256) self.speech_window = np.hamming(2 * self.source_cache_len) + # Must cover decode()'s own causal receptive field, not just the F0 + # margin; window_len=48 already passes + # test_incremental_hift_bounded_window_is_close, so 64 has headroom. + self._hift_window_len = 64 @property def input_frame_rate(self) -> int: @@ -226,40 +243,79 @@ class CosyVoice3Code2Wav(nn.Module): chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) cached_mel = None if not cache_state else cache_state.get("mel") - speech_offset_obj = None if not cache_state else cache_state.get("speech_offset") - try: - speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0 - except (TypeError, ValueError): - speech_offset = 0 - + cached_offset = 0 if not cache_state else int(cache_state.get("mel_offset", 0)) + window_len = self._hift_window_len + samples_per_mel = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"]) + trim = int(self.hift.f0_predictor.condnet[0].causal_padding) + f0_margin_frames = int(self.hift.f0_predictor.left_context_frames) + f0_margin = 0 if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0: cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype) - tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel + total_hist = cached_mel.shape[-1] + overlap = min(window_len + trim, total_hist) + window_mel = torch.cat([cached_mel[..., -overlap:], chunk_mel], dim=-1) + window_offset = cached_offset + total_hist - overlap + uv_offset = window_offset * samples_per_mel + f0_margin = min(f0_margin_frames, total_hist - overlap) + f0_input_mel = ( + torch.cat( + [cached_mel[..., total_hist - overlap - f0_margin : total_hist - overlap], window_mel], dim=-1 + ) + if f0_margin > 0 + else window_mel + ) else: - tts_mel = chunk_mel + window_mel = chunk_mel + f0_input_mel = chunk_mel + window_offset = 0 + uv_offset = 0 - if tts_mel.shape[-1] == 0: + phase_acc = None if not cache_state else cache_state.get("phase_acc") + if phase_acc is not None: + phase_acc = phase_acc.to(device=chunk_mel.device, dtype=chunk_mel.dtype) + + if window_mel.shape[-1] == 0: tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype) + new_phase_acc = phase_acc else: - tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize) + next_overlap = min(window_len + trim, window_mel.shape[-1]) + # On finalize the F0 predictor consumes the whole window (it does not + # hold back `trim` frames), so trim must not inflate the phase-carry + # index -- it would point past the end of a short window. + carry_trim = 0 if finalize else trim + tts_speech, _, new_phase_acc = self.hift.inference( + speech_feat=f0_input_mel, + finalize=finalize, + phase_acc=phase_acc, + uv_offset=uv_offset, + next_overlap=next_overlap, + trim=carry_trim, + f0_margin=f0_margin, + ) tts_speech = tts_speech.reshape(tts_speech.shape[0], -1) - speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1]))) - emitted_speech = tts_speech[:, speech_offset:] + new_samples = chunk_mel.shape[-1] * samples_per_mel + if finalize: + frames_withheld_per_chunk = trim + self.hift.conv_pre_look_right + 1 + released = frames_withheld_per_chunk * samples_per_mel + emitted_speech = tts_speech[:, -(new_samples + released) :] + else: + emitted_speech = tts_speech[:, -new_samples:] if new_samples > 0 else tts_speech[:, :0] if finalize: return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None new_state = { - "mel": tts_mel.detach().cpu().contiguous(), - "speech_offset": int(tts_speech.shape[-1]), + "mel": f0_input_mel.detach().cpu().contiguous(), + "mel_offset": window_offset - f0_margin, + "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None, } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state @torch.inference_mode() def forward_streaming_batch( self, - items: list[dict[str, object]], + items: list[StreamingFlowItem], *, n_timesteps: int = 10, ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: @@ -270,20 +326,16 @@ class CosyVoice3Code2Wav(nn.Module): group and passed to the flow as per-row token lengths. """ results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) - groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {} + groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {} for index, item in enumerate(items): - token = item["token"] - prompt_token = item["prompt_token"] - prompt_feat = item["prompt_feat"] - embedding = item["embedding"] - assert isinstance(token, torch.Tensor) - assert isinstance(prompt_token, torch.Tensor) - assert isinstance(prompt_feat, torch.Tensor) - assert isinstance(embedding, torch.Tensor) + assert isinstance(item["token"], torch.Tensor) + assert isinstance(item["prompt_token"], torch.Tensor) + assert isinstance(item["prompt_feat"], torch.Tensor) + assert isinstance(item["embedding"], torch.Tensor) key = ( - int(prompt_token.shape[1]), - int(prompt_feat.shape[1]), - int(embedding.shape[1]), + int(item["prompt_token"].shape[1]), + int(item["prompt_feat"].shape[1]), + int(item["embedding"].shape[1]), bool(item.get("finalize", False)), ) groups.setdefault(key, []).append((index, item)) @@ -305,11 +357,11 @@ class CosyVoice3Code2Wav(nn.Module): if len(group) == 1: index, item = group[0] result = self.forward_streaming( - token=item["token"], # type: ignore[arg-type] - prompt_token=item["prompt_token"], # type: ignore[arg-type] - prompt_feat=item["prompt_feat"], # type: ignore[arg-type] - embedding=item["embedding"], # type: ignore[arg-type] - cache_state=item.get("cache_state"), # type: ignore[arg-type] + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + cache_state=item.get("cache_state"), n_timesteps=n_timesteps, token_offset_tokens=int(item.get("token_offset_tokens", 0)), finalize=bool(item.get("finalize", False)), @@ -318,15 +370,13 @@ class CosyVoice3Code2Wav(nn.Module): continue token_tensors = [item["token"] for _, item in group] - assert all(isinstance(token, torch.Tensor) for token in token_tensors) token_lens = torch.tensor( - [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr] + [int(token.shape[1]) for token in token_tensors], dtype=torch.int32, ) max_token_len = int(token_lens.max().item()) padded_tokens = [] for token in token_tensors: - assert isinstance(token, torch.Tensor) if int(token.shape[1]) == max_token_len: padded_tokens.append(token) else: @@ -337,9 +387,9 @@ class CosyVoice3Code2Wav(nn.Module): ) padded_tokens.append(torch.cat([token, pad], dim=1)) tokens = torch.cat(padded_tokens, dim=0) - prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item] - prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item] - embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) + prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) + embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) finalize = bool(group[0][1].get("finalize", False)) @@ -370,7 +420,7 @@ class CosyVoice3Code2Wav(nn.Module): row_feat = row_feat[:, :, trim_mel:] results[index] = self._stream_hift_from_feat( row_feat, - cache_state=item.get("cache_state"), # type: ignore[arg-type] + cache_state=item.get("cache_state"), finalize=finalize, ) @@ -443,7 +493,7 @@ class CosyVoice3Code2Wav(nn.Module): dtype=tts_mel.dtype, ) else: - tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True) + tts_speech, _, _ = self.hift.inference(speech_feat=tts_mel, finalize=True) return tts_speech From 1053f7b0bdf0a1ac18c09c536d08eefc32bc2995 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:39:01 +0800 Subject: [PATCH 02/99] [Bugfix] Preserve benchmark metric sample counts in results (#7624) Signed-off-by: psv666 <2693925048@qq.com> Co-authored-by: amy-why-3459 --- .../benchmarks/patch/test_benchmark_result.py | 74 +++++++++++++++++++ tests/dfx/perf/tests/test_runner_metadata.py | 16 ++-- vllm_omni/benchmarks/patch/patch.py | 8 ++ 3 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 tests/benchmarks/patch/test_benchmark_result.py diff --git a/tests/benchmarks/patch/test_benchmark_result.py b/tests/benchmarks/patch/test_benchmark_result.py new file mode 100644 index 000000000..1e9619fe3 --- /dev/null +++ b/tests/benchmarks/patch/test_benchmark_result.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Verify metric sample counts survive benchmark JSON serialization.""" + +import asyncio +import json +import time + +import pytest +from vllm.benchmarks.datasets import SampleRequest +from vllm.benchmarks.serve import TaskType + +from vllm_omni.benchmarks.patch import patch + +pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu] + + +@pytest.mark.parametrize("selected_metrics", [["ttft", "tpot", "itl"], []]) +@pytest.mark.parametrize("measured", [True, False]) +def test_benchmark_result_preserves_tpot_sample_count(monkeypatch, measured, selected_metrics): + async def request_func(**kwargs): + return patch.MixRequestFuncOutput( + success=True, + prompt_len=2, + output_tokens=3, + generated_text="hello world!", + start_time=time.perf_counter(), + ttft=0.01, + latency=0.03, + text_latency=0.03, + itl=[0.01, 0.01] if measured else [0.0, 0.0], + ) + + monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-result", request_func) + result = asyncio.run( + patch.benchmark( + task_type=TaskType.GENERATION, + endpoint_type="test-result", + api_url="http://unused/v1/chat/completions", + base_url="http://unused", + model_id="test-model", + model_name="test-model", + tokenizer=None, + input_requests=[SampleRequest(prompt="hi", prompt_len=2, expected_output_len=3)], + logprobs=None, + request_rate=float("inf"), + burstiness=1.0, + disable_tqdm=True, + num_warmups=0, + profile=False, + selected_percentile_metrics=selected_metrics, + selected_percentiles=[50.0], + ignore_eos=False, + goodput_config_dict={}, + max_concurrency=1, + lora_modules=None, + extra_headers=None, + extra_body=None, + ready_check_timeout_sec=0, + ) + ) + # DFX reads this dictionary from the benchmark's JSON result file. + saved_result = json.loads(json.dumps(result)) + assert saved_result["completed"] == 1 + assert saved_result.get("num_tpot_samples") == int(measured) + assert saved_result["num_ttft_samples"] == 1 + assert saved_result["num_itl_samples"] == 2 + assert "num_audio_ttfp_samples" in saved_result + assert "num_audio_rtf_samples" in saved_result + if measured and "tpot" in selected_metrics: + assert saved_result["mean_tpot_ms"] == pytest.approx(10.0) + else: + assert "mean_tpot_ms" not in saved_result diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index df99f090d..a0f966121 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/tests/dfx/perf/tests/test_runner_metadata.py @@ -570,6 +570,7 @@ def test_omni_tpot_baseline_accepts_measured_finite_sample(): @pytest.mark.parametrize( ("num_tpot_samples", "mean_tpot_ms", "match"), [ + (None, 10.0, "no measurable TPOT samples"), (0, float("nan"), "no measurable TPOT samples"), (1, float("nan"), "mean_tpot_ms is not finite"), ], @@ -577,14 +578,17 @@ def test_omni_tpot_baseline_accepts_measured_finite_sample(): def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match): from tests.dfx.perf.scripts.run_benchmark import assert_result + result = { + "completed": 1, + "Hardware": "H100", + "mean_tpot_ms": mean_tpot_ms, + } + if num_tpot_samples is not None: + result["num_tpot_samples"] = num_tpot_samples + with pytest.raises(AssertionError, match=match): assert_result( - { - "completed": 1, - "Hardware": "H100", - "num_tpot_samples": num_tpot_samples, - "mean_tpot_ms": mean_tpot_ms, - }, + result, {"baseline": {"H100": {"mean_tpot_ms": 20.0}}}, 1, ) diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index bee910c90..5a5e74b0d 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -3078,6 +3078,14 @@ async def benchmark( "max_concurrent_requests": metrics.max_concurrent_requests, "rtfx": metrics.rtfx, } + for sample_count in ( + "num_ttft_samples", + "num_tpot_samples", + "num_itl_samples", + "num_audio_ttfp_samples", + "num_audio_rtf_samples", + ): + result[sample_count] = getattr(metrics, sample_count) else: result = { "duration": benchmark_duration, From 81b5be274d88eeb73e84091ba3f1466dc6f5a184 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Wed, 16 Sep 2026 23:54:28 +0800 Subject: [PATCH 03/99] [Bugfix][Qwen-Image] Restore RotaryEmbedding CUDA RoPE for Diffusers e2e (#7513) Signed-off-by: NumberWan Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> --- .../benchmark_qwen_image_qk_norm_rope.py | 12 +++---- .../test_qwen_image_fused_qk_norm_rope.py | 34 +++++++++++-------- .../qwen_image/test_qwen_image_rope_utils.py | 25 -------------- tests/e2e/accuracy/test_qwen_image.py | 6 ++-- .../qwen_image/qwen_image_transformer.py | 22 ++---------- 5 files changed, 31 insertions(+), 68 deletions(-) diff --git a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py index ee69f7e59..2eaa097bd 100644 --- a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py +++ b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py @@ -22,9 +22,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( _fused_cuda_supported, fused_qk_norm_rope, ) -from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( - _apply_qwen_image_rotary_emb, -) +from vllm_omni.diffusion.layers.rope import RotaryEmbedding def _parse_args() -> argparse.Namespace: @@ -87,10 +85,10 @@ def _native( ) -> tuple[torch.Tensor, torch.Tensor]: q = F.rms_norm(q, (q.shape[-1],), q_weight, eps) k = F.rms_norm(k, (k.shape[-1],), k_weight, eps) - return ( - _apply_qwen_image_rotary_emb(q, freqs), - _apply_qwen_image_rotary_emb(k, freqs), - ) + rope = RotaryEmbedding(is_neox_style=False) + cos = freqs.real.to(q.dtype) + sin = freqs.imag.to(q.dtype) + return rope(q, cos, sin), rope(k, cos, sin) def _fused( diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py index 2f10ea33e..16dc85115 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py @@ -11,7 +11,6 @@ import torch.nn as nn from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( - _apply_qwen_image_rotary_emb, _qwen_image_qk_norm_rope, ) @@ -73,22 +72,28 @@ def _make_input( return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs) -def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: +def _eager_rotary(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: q = data.norm_q(data.q) k = data.norm_k(data.k) - if data.q.device.type == "cuda": - return ( - _apply_qwen_image_rotary_emb(q, data.freqs), - _apply_qwen_image_rotary_emb(k, data.freqs), - ) - rope = RotaryEmbedding(is_neox_style=False) cos = data.freqs.real.to(data.q.dtype) sin = data.freqs.imag.to(data.q.dtype) return rope(q, cos, sin), rope(k, cos, sin) -def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: +def _fp32_complex_rotary(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Local copy of the old CUDA helper; fused kernel still tracks this math.""" + paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) + + +def _fused_kernel_reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: + q = data.norm_q(data.q) + k = data.norm_k(data.k) + return _fp32_complex_rotary(q, data.freqs), _fp32_complex_rotary(k, data.freqs) + + +def _run(data: QwenImageQKInput, *, use_fused: bool = True) -> tuple[torch.Tensor, torch.Tensor]: return _qwen_image_qk_norm_rope( data.q, data.k, @@ -97,6 +102,7 @@ def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: data.freqs, RotaryEmbedding(is_neox_style=False), EPS, + use_fused=use_fused, ) @@ -110,8 +116,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference(): packed_qkv_view=True, ) - actual_q, actual_k = _run(data) - expected_q, expected_k = _reference(data) + actual_q, actual_k = _run(data, use_fused=False) + expected_q, expected_k = _eager_rotary(data) torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5) torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5) @@ -127,8 +133,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference(): packed_qkv_view=True, ) - actual_q, actual_k = _run(data) - expected_q, expected_k = _reference(data) + actual_q, actual_k = _run(data, use_fused=False) + expected_q, expected_k = _eager_rotary(data) torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3) torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3) @@ -150,7 +156,7 @@ def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference( ) actual_q, actual_k = _run(data) - expected_q, expected_k = _reference(data) + expected_q, expected_k = _fused_kernel_reference(data) torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py index 4f3a81491..3fd141dc1 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py @@ -3,37 +3,12 @@ import torch from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenEmbedRope, - _apply_qwen_image_rotary_emb, ) from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("noncontiguous", [False, True]) -def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous): - from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen - - generator = torch.Generator().manual_seed(42) - x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype) - if noncontiguous: - x = x.transpose(1, 2).contiguous().transpose(1, 2) - angles = torch.randn(17, 64, generator=generator) - freqs = torch.polar(torch.ones_like(angles), angles) - expected = apply_rotary_emb_qwen(x, freqs, use_real=False) - actual = _apply_qwen_image_rotary_emb(x, freqs) - assert actual.dtype == dtype - assert actual.shape == x.shape - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - if dtype == torch.bfloat16: - # A regression to early BF16 frequency rounding must be observable. - rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float()) - rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs) - assert not torch.equal(actual, rounded) - - def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count(): prompt_embeds = torch.zeros(2, 64, 8) prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index fb2f8cbc8..1fca04f99 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -32,9 +32,9 @@ HEIGHT = 512 NUM_INFERENCE_STEPS = 20 TRUE_CFG_SCALE = 4.0 SEED = 42 -SSIM_THRESHOLD = 0.94 -# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8 -# vs Diffusers on the FA3 hub path (see #5734 / #5963 / #6971). +# H100 nightly after RotaryEmbedding CUDA RoPE: SSIM 0.978375 / PSNR 31.83 +# vs Diffusers (Buildkite #15339). SSIM 0.97 so a drop is visible. +SSIM_THRESHOLD = 0.97 PSNR_THRESHOLD = 27.0 MODEL_2512_ID = "Qwen/Qwen-Image-2512" diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 11568f62c..cb1adbb61 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -60,17 +60,6 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) -def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: - """Rotate interleaved pairs before rounding back to the activation dtype. - - Qwen-Image's reference uses complex FP32 multiplication. Rounding the - frequencies to BF16 before rotation loses positional precision; those - errors accumulate across the denoising steps. - """ - paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) - - def _qwen_image_qk_norm_rope( q: torch.Tensor, k: torch.Tensor, @@ -103,16 +92,11 @@ def _qwen_image_qk_norm_rope( fused_k.reshape(batch, seq_len, num_kv_heads, head_dim), ) + # Eager path: BF16/activation-dtype RotaryEmbedding on every device. + # CUDA used to call FP32 complex multiply here; that matches the Diffusers + # helper in unit tests but drops Omni vs Diffusers pipeline PSNR (#7494). q = norm_q(q) k = norm_k(k) - if q.device.type == "cuda": - return ( - _apply_qwen_image_rotary_emb(q, freqs), - _apply_qwen_image_rotary_emb(k, freqs), - ) - - # Retain the platform-specific kernels on other accelerators, which may - # not support complex tensors. cos = torch.real(freqs).to(q.dtype) sin = torch.imag(freqs).to(q.dtype) return rope(q, cos, sin), rope(k, cos, sin) From 3d75254a4401b9f146d22d5680c772251b8df82f Mon Sep 17 00:00:00 2001 From: y-null Date: Thu, 17 Sep 2026 00:00:45 +0800 Subject: [PATCH 04/99] [CI][MiniCPM-o] Lower the Daily-Omni accuracy gate to 0.77 (#7657) Signed-off-by: y_null --- tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py index ad659541d..acf717e42 100644 --- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py +++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py @@ -3,7 +3,7 @@ """MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage. Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches -~78% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``, +~77% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``, text modalities, and server ``--interleave-mm-strings`` + 1fps / 128-frame media-io kwargs, also pinned in ``minicpmo_4_5.yaml``). """ @@ -48,7 +48,12 @@ _RESULT_DIR = Path( ) ) -_MIN_DAILY_OMNI_ACCURACY = 0.78 +# The 0.78 gate (934/1197) sits 1 item above the observed NPU mean +# (~933/1197) and inside the observed 932~940 spread, so it flips nightly - +# 3 failures in the latest 7 dated NPU nightlies (#6887). 0.77 (922/1197) +# sits 10 items below the observed minimum (932) while still failing a +# >=1pp real accuracy regression. +_MIN_DAILY_OMNI_ACCURACY = 0.77 _MAX_SEED_TTS_MEAN_WER = 0.05 # Match the validated Daily-Omni client body from daily_omni_bench.sh. _DAILY_EXTRA_BODY = { From 2a6847171516b5977f4e9e835f427f5e731d73f3 Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Thu, 17 Sep 2026 01:09:23 +0800 Subject: [PATCH 05/99] [Bugfix] Reject failed diffusion benchmark warmups (#7092) Signed-off-by: hyw Signed-off-by: HuangYuwei Signed-off-by: Alex Brooks Co-authored-by: Alex Brooks --- .../diffusion/diffusion_benchmark_serving.py | 12 +- .../test_diffusion_benchmark_warmups.py | 128 ++++++++++++++++++ 2 files changed, 139 insertions(+), 1 deletion(-) create mode 100644 tests/benchmarks/test_diffusion_benchmark_warmups.py diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 1e5aaf29c..6bf64ec1d 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -1087,6 +1087,12 @@ async def _run_warmups( warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests] warmup_outputs = await asyncio.gather(*warmup_tasks) + failed_count = sum(not output.success for output in warmup_outputs) + if failed_count: + raise RuntimeError( + f"{failed_count}/{len(warmup_outputs)} warmup requests failed; " + "check server logs before running the benchmark." + ) return list(zip(warmup_requests, warmup_outputs)) @@ -1455,7 +1461,7 @@ async def benchmark(args): print(f"Metrics saved to {args.output_file}") -if __name__ == "__main__": +def build_parser(): parser = argparse.ArgumentParser(description="Benchmark serving for diffusion models.") parser.add_argument( "--base-url", @@ -1637,6 +1643,10 @@ if __name__ == "__main__": action="store_true", help="Request stage duration metrics from endpoints that support return_stage_metrics.", ) + return parser + +if __name__ == "__main__": + parser = build_parser() args = parser.parse_args() asyncio.run(benchmark(args)) diff --git a/tests/benchmarks/test_diffusion_benchmark_warmups.py b/tests/benchmarks/test_diffusion_benchmark_warmups.py new file mode 100644 index 000000000..9d72189b7 --- /dev/null +++ b/tests/benchmarks/test_diffusion_benchmark_warmups.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib +import sys + +import pytest + +from benchmarks.diffusion import backends +from benchmarks.diffusion.backends import RequestFuncInput, RequestFuncOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu, pytest.mark.asyncio] + +ENDPOINT = "/v1/chat/completions" +TEST_MODEL = "test-model" + + +@pytest.fixture +def diffusion_benchmark(monkeypatch): + # The standalone script imports its sibling backend module by name. + monkeypatch.setitem(sys.modules, "backends", backends) + return importlib.import_module("benchmarks.diffusion.diffusion_benchmark_serving") + + +@pytest.fixture +def warmup_args(diffusion_benchmark): + from benchmarks.diffusion.diffusion_benchmark_serving import build_parser + + parser = build_parser() + return parser.parse_args( + [ + f"--endpoint={ENDPOINT}", + "--dataset=random", + "--task=t2i", + "--warmup-requests=2", + "--warmup-num-inference-steps=4", + ] + ) + + +@pytest.fixture +def input_requests(): + return [ + RequestFuncInput( + prompt=prompt, + api_url=f"http://test.local{ENDPOINT}", + model=TEST_MODEL, + num_inference_steps=20, + ) + for prompt in ("first", "second") + ] + + +async def test_disabled_warmup_sends_no_requests(diffusion_benchmark, warmup_args, input_requests): + warmup_args.warmup_requests = 0 + + async def unexpected_request(*args): + pytest.fail("disabled warmup must not send requests") + + pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, unexpected_request) + + assert pairs == [] + + +async def test_successful_warmups_preserve_request_output_pairs(diffusion_benchmark, warmup_args, input_requests): + outputs = [RequestFuncOutput(success=True, latency=0.5), RequestFuncOutput(success=True, latency=1.0)] + pending_outputs = iter(outputs) + + async def successful_request(req, session, pbar): + return next(pending_outputs) + + pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, successful_request) + + assert [request.prompt for request, _ in pairs] == ["first", "second"] + assert [request.num_inference_steps for request, _ in pairs] == [4, 4] + assert [output for _, output in pairs] == outputs + assert [request.num_inference_steps for request in input_requests] == [20, 20] + + +@pytest.mark.parametrize("failed_count", [1, 2], ids=["partial-failure", "all-failed"]) +@pytest.mark.parametrize( + "error", + ["HTTP 503: private prompt rejected", "Cannot connect to https://test.local/?token=private"], + ids=["http-error", "transport-error"], +) +async def test_failed_warmups_abort_without_exposing_response_errors( + diffusion_benchmark, warmup_args, input_requests, failed_count, error +): + outputs = iter( + [RequestFuncOutput(success=False, error=error) for _ in range(failed_count)] + + [RequestFuncOutput(success=True) for _ in range(2 - failed_count)] + ) + + async def request(req, session, pbar): + return next(outputs) + + with pytest.raises(RuntimeError, match=f"{failed_count}/2 warmup requests failed") as exc_info: + await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, request) + + assert error not in str(exc_info.value) + assert "server logs" in str(exc_info.value) + + +async def test_benchmark_does_not_start_measurement_after_failed_warmup( + monkeypatch, mocker, diffusion_benchmark, warmup_args, input_requests +): + sent = [] + + async def failed_request(req, session, pbar): + sent.append(req) + return RequestFuncOutput(success=False, error="HTTP 503: unavailable") + + monkeypatch.setitem( + diffusion_benchmark.backends_function_mapping["2i"], + warmup_args.endpoint, + (failed_request, warmup_args.endpoint), + ) + dataset = mocker.Mock(spec=diffusion_benchmark.RandomDataset) + dataset.get_requests.return_value = input_requests + monkeypatch.setattr(diffusion_benchmark, "RandomDataset", mocker.Mock(return_value=dataset)) + timer = mocker.Mock(spec=diffusion_benchmark.time) + timer.perf_counter.side_effect = AssertionError("measurement started after failed warmup") + monkeypatch.setattr(diffusion_benchmark, "time", timer) + + with pytest.raises(RuntimeError, match="2/2 warmup requests failed"): + await diffusion_benchmark.benchmark(warmup_args) + + assert len(sent) == warmup_args.warmup_requests From d4dba0ddb1d645b8e546a772f3b08e97835ac56c Mon Sep 17 00:00:00 2001 From: THUqliu <59546143+THUqliu@users.noreply.github.com> Date: Thu, 17 Sep 2026 01:13:32 +0800 Subject: [PATCH 06/99] [Misc] Remove dead deploy keys from personaplex.yaml (#7454) Signed-off-by: THUqliu <3298452736@qq.com> --- vllm_omni/deploy/personaplex.yaml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm_omni/deploy/personaplex.yaml b/vllm_omni/deploy/personaplex.yaml index 9f1da7e63..c4cc472d1 100644 --- a/vllm_omni/deploy/personaplex.yaml +++ b/vllm_omni/deploy/personaplex.yaml @@ -25,12 +25,8 @@ connectors: extra: shm_threshold_bytes: 65536 codec_streaming: true - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 # Mimi runs at 12.5 Hz (80 ms / frame). Stream a few frames per chunk. codec_chunk_frames: 5 - codec_left_context_frames: 10 # Emit only the first audio chunk early, then return to codec_chunk_frames. initial_codec_chunk_frames: 1 From 4392af5ce37ea346a3b51914719a2679dd59d416 Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Wed, 16 Sep 2026 11:36:52 -0700 Subject: [PATCH 07/99] [Bugfix][Qwen3-Omni] Include deferred residual in captured Thinker states (#7304) Signed-off-by: Deep Shah --- .../test_qwen3_omni_forward_contract.py | 51 ++++++++++++++++++- .../qwen3_omni/qwen3_omni_moe_thinker.py | 5 +- 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py index 0055d1704..f4635a81f 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for the Qwen3-Omni thinker forward return contract. Background @@ -41,6 +41,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( Qwen3OmniMoeForConditionalGeneration, ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3MoeLLMModel, Qwen3OmniMoeThinkerForConditionalGeneration, ) @@ -54,6 +55,52 @@ def _hidden_states() -> torch.Tensor: return torch.zeros(_TOKENS, _HIDDEN) +def test_single_rank_intermediate_capture_includes_deferred_residual(mocker): + """Capture the Talker state on the default single-PP-rank path.""" + import vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker as module + + pp_group = mocker.Mock(is_first_rank=True, is_last_rank=True) + mocker.patch.object(module, "get_pp_group", return_value=pp_group) + + class SplitStateLayer(nn.Module): + def forward(self, positions, hidden_states, residual): + next_residual = hidden_states if residual is None else hidden_states + residual + return hidden_states * 2, next_residual + + class Combine(nn.Module): + def forward(self, hidden_states, residual): + return hidden_states + residual, None + + model = object.__new__(Qwen3MoeLLMModel) + nn.Module.__init__(model) + model.layers = nn.ModuleList([SplitStateLayer(), SplitStateLayer()]) + model.start_layer = 0 + model.end_layer = 2 + model.norm = Combine() + inputs = torch.ones(_TOKENS, _HIDDEN) + + output_without_capture, _ = Qwen3MoeLLMModel.forward( + model, + input_ids=None, + positions=torch.arange(_TOKENS), + inputs_embeds=inputs, + ) + output_with_capture, captured = Qwen3MoeLLMModel.forward( + model, + input_ids=None, + positions=torch.arange(_TOKENS), + inputs_embeds=inputs, + capture_layer_indices=[0, 1], + return_hidden_states=True, + ) + + layers = captured["hidden_states"]["layers"] + torch.testing.assert_close(layers[0], inputs) + torch.testing.assert_close(layers[1], torch.full_like(inputs, 3)) + torch.testing.assert_close(output_with_capture, output_without_capture) + torch.testing.assert_close(output_with_capture, torch.full_like(inputs, 7)) + + class _InnerModelStub: """Stands in for ``Qwen3MoeLLMModel``: always returns the internal 2-tuple.""" @@ -169,7 +216,7 @@ def test_make_omni_output_accepts_bare_tensor(): def test_make_omni_output_accepts_capture_tuple(): model = _make_combined(is_staged_run=True) - captured = {"hidden_states": {"layers": {}}} + captured: dict[str, dict[str, dict[int, torch.Tensor]]] = {"hidden_states": {"layers": {}}} out = model.make_omni_output((_hidden_states(), captured)) assert isinstance(out, OmniOutput) assert out.multimodal_outputs == captured diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 8d0389d63..c78679c55 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -588,7 +588,10 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): if layer_idx in capture_set: hs = captured_hidden_states.setdefault("hidden_states", {}) layers = hs.setdefault("layers", {}) - layers[layer_idx] = hidden_states.clone().view(-1, hidden_states.shape[-1]) + # vLLM defers the residual addition until the next RMSNorm. + # Reconstruct the logical decoder state before capturing it. + captured = hidden_states.clone() if residual is None else hidden_states + residual + layers[layer_idx] = captured.view(-1, captured.shape[-1]) hidden_states, residual = layer( positions, From 38c7eeee4c7662d28db0c8a3762ebce98d3a520a Mon Sep 17 00:00:00 2001 From: Lei Ke <1141466880@qq.com> Date: Thu, 17 Sep 2026 02:38:06 +0800 Subject: [PATCH 08/99] [BugFix][NPU] Fix Fish Speech S2 Pro NPU support (#7546) Signed-off-by: Lei Ke <1141466880@qq.com> --- tests/attention/test_fish_kvcache_attn.py | 10 ++++++++++ tests/config/test_config_factory.py | 7 +++++++ vllm_omni/attention/fish_kvcache_attn.py | 8 ++++++++ vllm_omni/deploy/fish_qwen3_omni.yaml | 9 +++++++++ .../models/fish_speech/fish_speech_slow_ar.py | 7 ++++++- 5 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py index 4a8ea17e8..73d44295a 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.py @@ -3,6 +3,7 @@ import pytest import torch +from vllm.platforms import current_platform from vllm_omni.attention import fish_kvcache_attn, fish_kvcache_backend @@ -113,6 +114,15 @@ def test_fish_kvcache_enabled_by_default(monkeypatch): assert not fish_kvcache_attn.is_fish_kvcache_attn_required() +def test_fish_kvcache_fast_path_is_unavailable_off_cuda(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + monkeypatch.setattr( + fish_kvcache_attn, "_triton_backend", lambda: pytest.fail("Triton backend must not load off CUDA") + ) + + assert not fish_kvcache_attn.is_available() + + @pytest.mark.parametrize("value", ["", "1", "true", "yes", "on", "required"]) def test_fish_kvcache_enabled_values_are_consistent(monkeypatch, value): monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index c81a4b3b7..7940e2176 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2688,6 +2688,13 @@ class TestPlatformOverrides: # Explicit null clears the inherited single-GPU 2 GiB CUDA cap. assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None + def test_fish_speech_npu_uses_ascend_kv_block_size(self): + deploy_path = Path(get_deploy_config_path("fish_qwen3_omni.yaml")) + + deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu") + + assert deploy.stages[0].engine_extras["block_size"] == 128 + def test_npu_overrides(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): diff --git a/vllm_omni/attention/fish_kvcache_attn.py b/vllm_omni/attention/fish_kvcache_attn.py index e7235168c..6e402ca54 100644 --- a/vllm_omni/attention/fish_kvcache_attn.py +++ b/vllm_omni/attention/fish_kvcache_attn.py @@ -24,6 +24,14 @@ def _triton_backend() -> Any: def is_available() -> bool: + # The Fish kvcache fast path is CUDA-only: it depends on FlashAttention's + # use_cascade metadata (absent from Ascend's metadata) and a Triton CUDA + # kernel. On other platforms report unavailable so the native attention + # backend takes over instead of crashing on attn_metadata.use_cascade. + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + return False return _triton_backend().is_available() diff --git a/vllm_omni/deploy/fish_qwen3_omni.yaml b/vllm_omni/deploy/fish_qwen3_omni.yaml index f10ed1656..0f4866a1f 100644 --- a/vllm_omni/deploy/fish_qwen3_omni.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni.yaml @@ -64,6 +64,15 @@ stages: repetition_penalty: 1.0 platforms: + npu: + # Ascend paged attention only accepts 128-token cache blocks + # (AscendAttentionBackend.get_supported_kernel_block_sizes() == [128]), + # while vLLM's CacheConfig default is 16. Without this override stage 0 + # aborts during KV-cache init with: + # ValueError: No common block size for 16. + stages: + - stage_id: 0 + block_size: 128 xpu: stages: # Verified on Intel Arc B60. diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py index 1317dc044..57f37f3d9 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py @@ -32,6 +32,7 @@ from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.utils.speaker_cache import get_speaker_cache from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig @@ -195,7 +196,11 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module): self.mtp_hidden_size = int(self.text_config.hidden_size) self.talker_mtp_output_key = ("codes", "audio") self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")} - self.talker_mtp_graph_safe = True + # NPU (Ascend) ACL graphs cannot capture the Fast-AR random-sampling ops + # (torch.nonzero / torch.multinomial used for top-p), so keep the talker + # MTP eager on NPU to fall back to native (non-graph) execution. Mirrors + # the 310P patch for Qwen3-TTS (talker_mtp_graph_safe=False). + self.talker_mtp_graph_safe = not current_omni_platform.is_npu() # Qwen3 transformer backbone. self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) From c97b5fc741ea44acaed37efe520ef8600d1c2abf Mon Sep 17 00:00:00 2001 From: shiyichuan <93317314+CarrotSwordsman@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:38:45 +0800 Subject: [PATCH 09/99] [Bugfix][Frontend] Route synthetic aborts through final stage (#7006) Signed-off-by: mershi Co-authored-by: mershi --- tests/entrypoints/test_omni_entrypoints.py | 38 +++++++++++++++++++ vllm_omni/entrypoints/async_omni.py | 1 + vllm_omni/entrypoints/async_omni_base.py | 30 ++++++++++++--- vllm_omni/entrypoints/client_request_state.py | 5 +++ 4 files changed, 69 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 32d7d9988..ea57d07ab 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import gc import queue from collections.abc import Callable @@ -634,6 +635,43 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc assert "req-1" not in app.request_states +@pytest.mark.asyncio +async def test_async_omni_abort_yields_terminal_from_requested_final_stage(monkeypatch: pytest.MonkeyPatch): + stage_metadata = [ + _stage_meta(stage_type="llm", final_output=False, final_output_type=None), + _stage_meta(stage_type="llm", final_output=True, final_output_type="audio"), + ] + engine = FakeAsyncOmniEngine(stage_metadata=stage_metadata) + _patch_engine(monkeypatch, engine) + app = AsyncOmni("dummy-model") + + async def collect_outputs() -> list[OmniRequestOutput]: + return [ + output + async for output in app.generate( + prompt="hello", + request_id="req-1", + output_modalities=["audio"], + ) + ] + + try: + generate_task = asyncio.create_task(collect_outputs()) + while not engine.submitted: + await asyncio.sleep(0) + await app.abort("req-1") + outputs = await asyncio.wait_for(generate_task, timeout=1) + finally: + app.shutdown() + + assert len(outputs) == 1 + assert outputs[0].stage_id == 1 + assert outputs[0].final_output_type == "audio" + assert outputs[0].finished is True + assert outputs[0].outputs[0].finish_reason == "abort" + assert "req-1" not in app.request_states + + @pytest.mark.asyncio async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6114bbea2..1291e68f0 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -236,6 +236,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): req_state = ClientRequestState( request_id=request_id, external_request_id=external_request_id, + final_stage_id=final_stage_id_for_e2e, ) req_state.metrics = metrics req_state.request_arrival_ts = wall_start_ts diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index e7c86de97..e9bf2c77e 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -479,7 +479,22 @@ class AsyncOmniBase(OmniBase): if state is not None and rid not in delivered: queue = getattr(state, "queue", None) if queue is not None: - await state.queue.put(self._synthetic_abort_output_message(rid)) + final_stage_id = getattr(state, "final_stage_id", None) + if final_stage_id is None: + final_stage_id = 0 + stage_metadata = getattr(self, "_stage_meta_list", ()) + final_output_type = ( + getattr(stage_metadata[final_stage_id], "final_output_type", None) + if final_stage_id < len(stage_metadata) + else None + ) or "text" + await state.queue.put( + self._synthetic_abort_output_message( + rid, + stage_id=final_stage_id, + final_output_type=final_output_type, + ) + ) delivered.add(rid) for rid in request_ids: self._record_request_failure_once(rid, reason="client_abort") @@ -491,13 +506,18 @@ class AsyncOmniBase(OmniBase): logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) @staticmethod - def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + def _synthetic_abort_output_message( + request_id: str, + *, + stage_id: int, + final_output_type: str, + ) -> OutputMessage: """Terminal abort OutputMessage used when the engine returned none.""" engine_output = OmniRequestOutput( request_id=request_id, finished=True, - stage_id=0, - final_output_type="text", + stage_id=stage_id, + final_output_type=final_output_type, outputs=[ CompletionOutput( index=0, @@ -512,7 +532,7 @@ class AsyncOmniBase(OmniBase): ) return OutputMessage( request_id=request_id, - stage_id=0, + stage_id=stage_id, replica_id=None, engine_outputs=engine_output, metrics=None, diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 22f1c6d5e..ca9e96bab 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio from vllm_omni.metrics import OrchestratorAggregator @@ -11,10 +14,12 @@ class ClientRequestState: request_id: str, external_request_id: str | None = None, queue: asyncio.Queue | None = None, + final_stage_id: int | None = None, ): self.request_id = request_id self.external_request_id = external_request_id self.stage_id: int | None = None + self.final_stage_id: int | None = final_stage_id self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None # Request-scoped idempotency guard for Prometheus failure counters. From 04c55d0164e586373a2d3a3fdd8ab792b72017e8 Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Thu, 17 Sep 2026 02:39:13 +0800 Subject: [PATCH 10/99] Preserve audio boundary embeddings in Qwen2.5-Omni interleaved input (#7510) Signed-off-by: nodeeeeee --- .../test_qwen2_5_omni_expansion.py | 26 ++++++++++ .../qwen2_5_omni/test_per_video_audio_mask.py | 2 +- .../qwen2_5_omni/test_qwen2_5_omni_embed.py | 50 +++++++++++++++++++ .../qwen2_5_omni/qwen2_5_omni_thinker.py | 8 ++- 4 files changed, 83 insertions(+), 3 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py index 1e5e59aed..389b0b6ba 100644 --- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py @@ -96,6 +96,32 @@ def test_mix_to_text_audio_001(omni_server, online_client) -> None: online_client.send_omni_request(request_config) +@pytest.mark.slow +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stream", [False, True]) +def test_audio_in_video_boundaries(omni_server, online_client, stream: bool) -> None: + """Interleaved video/audio input must not treat audio boundaries as image placeholders.""" + # More than two seconds exercises multiple interleaved video/audio chunks. + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + content_text="Describe what you see and hear in this video in one sentence.", + ) + responses = online_client.send_omni_request( + { + "model": omni_server.model, + "messages": messages, + "stream": stream, + "modalities": ["text"], + "use_audio_in_video": True, + } + ) + assert responses[0].text_content + + @pytest.mark.slow @pytest.mark.omni @pytest.mark.parametrize("omni_server", test_params, indirect=True) diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index 54efc99e4..9ebad492b 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -582,7 +582,7 @@ def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hit def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): - fake_self = _fake_processor() + fake_self = _fake_qwen2_prompt_processor() video_placeholders = [ PlaceholderFeaturesInfo( modality="video", diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py index 7b624ba11..0d7cd8175 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py +++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py @@ -331,6 +331,56 @@ class TestEmbedInputIds: "Interleaved: audio positions should get audio embeddings" ) + def test_interleaved_audio_boundaries_keep_text_embeddings(self, mocker: MockerFixture): + from vllm.multimodal.processing.processor import PlaceholderFeaturesInfo + + from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerMultiModalProcessor, + ) + + processor = mocker.Mock() + processor.info.get_hf_processor.return_value = mocker.Mock(audio_token="audio", video_token="video") + processor.info.get_tokenizer.return_value.get_vocab.return_value = { + "audio": AUDIO_TOKEN_ID, + "video": VIDEO_TOKEN_ID, + } + # Real audio_bos/audio_eos IDs from the failing use_audio_in_video input. + tokens = [151647, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, 151648] + placeholder = PlaceholderFeaturesInfo( + modality="video", + item_idx=0, + start_idx=2, + tokens=tokens, + is_embed=torch.ones(len(tokens), dtype=torch.bool), + ) + placeholders = Qwen2_5OmniThinkerMultiModalProcessor._derive_audio_from_video_placeholders( + processor, + {"video": [placeholder]}, + {"audio": [[object()]]}, + [True], + ) + + input_ids = torch.tensor([TEXT_TOKEN_ID] * 2 + tokens + [TEXT_TOKEN_ID]) + is_multimodal = torch.zeros_like(input_ids, dtype=torch.bool) + for modality in ("audio", "video"): + derived = placeholders[modality][0].to_range() + is_multimodal[derived.offset : derived.offset + derived.length] |= derived.is_embed + + model, hidden = make_mock_model(mocker, thinker_cls=_omni_qwen2_5_thinker_cls()) + text_embeddings = model.get_language_model().embed_input_ids(input_ids) + result = model.embed_input_ids( + input_ids, + build_mm_embeds(audio_n=2, image_n=0, video_n=2, hidden=hidden), + is_multimodal=is_multimodal, + ) + + is_audio = input_ids == AUDIO_TOKEN_ID + is_video = input_ids == VIDEO_TOKEN_ID + assert torch.equal(is_multimodal, is_audio | is_video) + torch.testing.assert_close(result[is_audio], torch.full((2, hidden), 10.0)) + torch.testing.assert_close(result[is_video], torch.full((2, hidden), 30.0)) + torch.testing.assert_close(result[~(is_audio | is_video)], text_embeddings[~(is_audio | is_video)]) + # --------------------------------------------------------------------------- # Tests for merge_interleaved_embeddings helper diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 263ffa2ab..b7b7d074c 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -758,6 +758,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( tokenizer = self.info.get_tokenizer() processor = self.info.get_hf_processor() audio_token_id = tokenizer.get_vocab()[processor.audio_token] + video_token_id = tokenizer.get_vocab()[processor.video_token] result_placeholders = dict(placeholders) audio_placeholders = [] @@ -765,7 +766,10 @@ class Qwen2_5OmniThinkerMultiModalProcessor( audio_idx = 0 for video_idx, video_placeholder in enumerate(placeholders["video"]): - audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + placeholder_tokens = torch.tensor(video_placeholder.tokens) + audio_is_embed = placeholder_tokens == audio_token_id + # Audio boundary tokens keep their text embeddings. + video_is_embed = placeholder_tokens == video_token_id if video_use_audio_in_video[video_idx]: audio_placeholder = PlaceholderFeaturesInfo( @@ -783,7 +787,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( item_idx=video_idx, start_idx=video_placeholder.start_idx, tokens=video_placeholder.tokens, - is_embed=~audio_is_embed, + is_embed=video_is_embed, ) video_placeholders.append(video_placeholder_with_mask) From a8434a60da1f40a0d118a326908fd20c9ebe75c7 Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Thu, 17 Sep 2026 00:09:51 +0530 Subject: [PATCH 11/99] [BugFix][TTS] Fix MOSS Realtime serving and generation (#5661) Signed-off-by: akshatvishu --- .../test_moss_tts_expansion.py | 9 +- .../test_moss_tts_realtime.py | 5 +- .../test_moss_tts_realtime_expansion.py | 5 +- .../online_serving/test_moss_tts_expansion.py | 28 +- .../test_moss_tts_realtime_serving.py | 302 ++++++++++++++++++ .../moss_tts/test_audio_tokenizer_v2.py | 54 ++++ .../moss_tts/test_realtime_chunked_prefill.py | 86 +++++ .../test_realtime_generation_limit.py | 53 +++ .../models/moss_tts/test_realtime_prompt.py | 78 +++++ .../openai/tts_adapters/moss_tts.py | 71 +++- .../moss_tts/modeling_moss_tts_talker.py | 57 ++-- .../models/moss_tts/realtime_prompt.py | 60 ++++ .../models/moss_tts/reference_encoder.py | 27 +- 13 files changed, 771 insertions(+), 64 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py create mode 100644 tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_generation_limit.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_prompt.py create mode 100644 vllm_omni/model_executor/models/moss_tts/realtime_prompt.py diff --git a/tests/e2e/offline_inference/test_moss_tts_expansion.py b/tests/e2e/offline_inference/test_moss_tts_expansion.py index 81a64fcd9..b899252b4 100644 --- a/tests/e2e/offline_inference/test_moss_tts_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B). Weekly expansion coverage (``pytest.mark.slow``) for the delay-model path. @@ -74,7 +74,6 @@ def _get_test_config() -> str: pytestmark = [ pytest.mark.slow, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.tts, pytest.mark.parametrize( "omni_runner", @@ -137,7 +136,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor return torch.cat(chunks, dim=0), sr_final -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None: """VoiceGenerator: English instruction produces non-empty 24 kHz audio.""" req = _build_request( @@ -151,7 +150,7 @@ def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None: assert not torch.all(audio == 0), "Audio is silence" -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None: """VoiceGenerator: Chinese input produces non-empty audio.""" req = _build_request( @@ -165,7 +164,7 @@ def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None: assert not torch.all(audio == 0) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_batch(omni_runner: OmniRunner) -> None: """VoiceGenerator: batch of two requests each returns non-empty audio.""" requests = [ diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py index 81d97cd74..0f0bf59c7 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -95,7 +95,6 @@ def _get_test_config() -> str: # --------------------------------------------------------------------------- pytestmark = [ - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize( @@ -272,7 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor # --------------------------------------------------------------------------- -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None: """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.") @@ -283,7 +282,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) assert not torch.all(audio == 0), "Audio is silence" -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None: """MossTTSRealtime: Chinese input produces non-empty audio.""" req = _build_request(ref_audio_path, "你好，这是语音合成测试。") diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py index b84e5a0ff..678fc7e55 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). Exercises the local-transformer generation path on L4. Delay-model coverage is in @@ -33,7 +33,6 @@ _OMNI_RUNNER_PARAM = ( pytestmark = [ pytest.mark.slow, pytest.mark.tts, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] @@ -118,7 +117,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor raise AssertionError("No stage outputs received") -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path) -> None: """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" req = _build_request("This is a real-time TTS streaming test.", ref_audio_path) diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py index 598326a1e..bd38a7fc7 100644 --- a/tests/e2e/online_serving/test_moss_tts_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -30,20 +30,16 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -_SKIP_ISSUE_6417 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/6417", -) - pytestmark = [ + pytest.mark.advanced_model, pytest.mark.slow, pytest.mark.tts, - _SKIP_ISSUE_6417, ] MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav" -# Voice-clone output is not reliably transcribed by the Whisper check used at -# full_model run_level; assert non-trivial WAV payload size instead. +# Reject a response that contains only a header or a very short audio payload. +# Advanced and full runs also validate the transcript. _MIN_AUDIO_BYTES = 10_000 @@ -77,7 +73,9 @@ def get_prompt(prompt_type: str = "text") -> str: """ prompts = { "text": "Hello, this is a real-time voice cloning demo for testing.", - "chinese": "你好，这是一段实时语音合成测试。", + # Match the offline test prompt. Whisper-small can transcribe "一段" + # as "一本", which makes this short sample fail the similarity check. + "chinese": "你好，这是语音合成测试。", } return prompts.get(prompt_type, prompts["text"]) @@ -94,7 +92,7 @@ tts_server_params = [ ] -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -105,9 +103,8 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No Input Setting: stream=False Datasets: single request - NOTE: ``min_audio_bytes`` skips Whisper transcript similarity — cloned - voice output is often empty or mismatched under ASR without indicating a - real serving regression. + ``min_audio_bytes`` adds a payload-size check. It does not disable the + transcript check used by advanced and full runs. """ request_config = { "model": omni_server.model, @@ -121,7 +118,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -150,7 +147,7 @@ def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_ online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -161,7 +158,8 @@ def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_ur Input Setting: stream=False Datasets: single request - NOTE: same ``min_audio_bytes`` rationale as test_text_to_audio_001. + ``min_audio_bytes`` adds a payload-size check. It does not disable the + transcript check used by advanced and full runs. """ request_config = { "model": omni_server.model, diff --git a/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py new file mode 100644 index 000000000..d69597634 --- /dev/null +++ b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.tts_adapters import moss_tts as adapter_module +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter +from vllm_omni.model_executor.models.moss_tts import reference_encoder +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("hf_config", "expected_codec_path"), + [ + (SimpleNamespace(), "OpenMOSS-Team/MOSS-Audio-Tokenizer"), + ( + SimpleNamespace(codec_model_name_or_path="/models/custom-moss-codec"), + "/models/custom-moss-codec", + ), + ], + ids=["default-codec", "configured-codec"], +) +def test_realtime_components_use_the_realtime_model_and_codec( + monkeypatch: pytest.MonkeyPatch, + hf_config: SimpleNamespace, + expected_codec_path: str, +) -> None: + engine_client = SimpleNamespace( + model_config=SimpleNamespace( + model="OpenMOSS-Team/MOSS-TTS-Realtime", + hf_config=hf_config, + ) + ) + + server = MossTTSAdapter(SpeechServingContext(server=object(), engine_client=engine_client)) + + tokenizer = object() + codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})() + processor_calls = [] + + class Processor: + def __init__(self, *, tokenizer): + processor_calls.append(tokenizer) + + class_calls = [] + tokenizer_calls = [] + codec_calls = [] + + def fake_get_class(class_reference, model_id): + class_calls.append((class_reference, model_id)) + return Processor + + def fake_load_tokenizer(model_id, *, trust_remote_code): + tokenizer_calls.append((model_id, trust_remote_code)) + return tokenizer + + def fake_load_codec(model_id, *, trust_remote_code): + codec_calls.append((model_id, trust_remote_code)) + return codec + + monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", fake_get_class) + monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", fake_load_tokenizer) + monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", fake_load_codec) + + components = server._get_moss_realtime_components() + + assert components[0] is tokenizer + assert components[2] is codec + assert server._get_moss_realtime_components() is components + assert class_calls == [ + ( + "processing_mossttsrealtime.MossTTSRealtimeProcessor", + "OpenMOSS-Team/MOSS-TTS-Realtime", + ) + ] + assert tokenizer_calls == [("OpenMOSS-Team/MOSS-TTS-Realtime", True)] + assert codec_calls == [(expected_codec_path, True)] + assert processor_calls == [tokenizer] + + +def test_realtime_serving_builds_the_talker_prompt(monkeypatch: pytest.MonkeyPatch) -> None: + server = MossTTSAdapter(SpeechServingContext(server=SimpleNamespace(uploaded_speakers={}))) + server._moss_variant = "realtime" + tokenizer = object() + processor = object() + codec = object() + request_thread = threading.get_ident() + + def load_components(): + assert threading.get_ident() != request_thread + return tokenizer, processor, codec + + server._get_moss_realtime_components = load_components + + reference_codes = torch.arange(64, dtype=torch.int64).reshape(4, 16) + encode = AsyncMock(return_value=([reference_codes], {0: "reference-cache-key"})) + monkeypatch.setattr(server, "_encode_moss_references", encode) + + build_call = None + + def fake_build(actual_tokenizer, actual_processor, text, actual_codes): + nonlocal build_call + build_call = (actual_tokenizer, actual_processor, text, actual_codes) + return { + "prompt_token_ids": [10, 11], + "codes": {"ref": actual_codes}, + "ids": {"all": [12]}, + } + + monkeypatch.setattr(adapter_module, "build_realtime_prompt", fake_build) + + params = asyncio.run( + server._build_moss_tts_params( + OpenAICreateSpeechRequest( + input="speak this text", + ref_audio="data:audio/wav;base64,AAAA", + max_new_tokens=50, + ) + ) + ) + + encode.assert_awaited_once() + assert encode.call_args.kwargs == {"has_inline_ref_audio": False, "two_speaker": False} + assert encode.call_args.args[0].ref_audio == "data:audio/wav;base64,AAAA" + assert build_call == (tokenizer, processor, "speak this text", reference_codes) + assert params == { + "prompt_token_ids": [10, 11], + "codes": {"ref": reference_codes}, + "ids": {"all": [12]}, + "max_new_frames": [50], + "ref_audio_cache_key": "reference-cache-key", + } + assert "prompt_audio_array" not in params + + +def test_cancelled_component_waiter_does_not_duplicate_codec_load(monkeypatch): + adapter = MossTTSAdapter( + SpeechServingContext( + server=object(), + engine_client=SimpleNamespace( + model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Realtime", hf_config=SimpleNamespace()) + ), + ) + ) + started = threading.Event() + release = threading.Event() + second_started = threading.Event() + loads = [] + codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})() + + def load_codec(*args, **kwargs): + loads.append(args) + started.set() + assert release.wait(5), "codec loader was not released" + return codec + + monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", lambda *args: lambda **kwargs: object()) + monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object()) + monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", load_codec) + + def second_load(): + second_started.set() + return adapter._get_moss_realtime_components() + + async def run(): + first = asyncio.create_task(asyncio.to_thread(adapter._get_moss_realtime_components)) + try: + assert await asyncio.to_thread(started.wait, 5) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + second = asyncio.create_task(asyncio.to_thread(second_load)) + try: + assert await asyncio.to_thread(second_started.wait, 5) + finally: + release.set() + components = await asyncio.wait_for(second, 5) + assert components[2] is codec + assert adapter._get_moss_realtime_components() is components + assert len(loads) == 1 + finally: + release.set() + + asyncio.run(run()) + + +@pytest.mark.parametrize("voice_name", [None, "speaker"]) +def test_realtime_reference_cache_reuses_codes(voice_name): + cache = SpeakerEmbeddingCache() + resolve = AsyncMock(return_value=([0.0], 24000, "audio-key")) + request_thread = threading.get_ident() + encode_calls = [] + expected = torch.arange(64).reshape(4, 16) + + def batch_encode(wavs, *, num_quantizers): + assert threading.get_ident() != request_thread + assert num_quantizers == 16 + encode_calls.append(wavs) + return SimpleNamespace(audio_codes=expected.T.unsqueeze(1), audio_codes_lengths=torch.tensor([4])) + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache + ) + + async def run(): + kwargs = dict( + resolve_ref_audio=resolve, + get_artifact_key=lambda key: "content-key", + voice_name=voice_name, + voice_created_at=1, + ) + try: + first, first_key = await encoder.encode("reference", **kwargs) + second, second_key = await encoder.encode("reference", **kwargs) + assert torch.equal(first, expected) + assert torch.equal(second, expected) + assert first.dtype == torch.int64 + assert first_key == "audio-key" + assert second_key == (None if voice_name else "audio-key") + second.zero_() + third, _ = await encoder.encode("another-reference", **kwargs) + assert torch.equal(third, expected) + assert len(encode_calls) == 1 + assert resolve.await_count == (1 if voice_name else 3) + key = cache.make_cache_key( + voice_name or "ref:content-key", "moss_tts_realtime_nq16", 1 if voice_name else 0 + ) + assert cache.get(key)["codes"].dtype == torch.int32 + if voice_name: + kwargs["voice_created_at"] = 2 + await encoder.encode("reference", **kwargs) + assert len(encode_calls) == 2 + assert resolve.await_count == 2 + finally: + await encoder.aclose() + + asyncio.run(run()) + + +def test_realtime_reference_batch_trims_each_item(): + cache = SpeakerEmbeddingCache() + codes = torch.arange(16 * 2 * 5).reshape(16, 2, 5) + calls = [] + + def batch_encode(wavs, *, num_quantizers): + calls.append(wavs) + assert num_quantizers == 16 + assert [wav.shape for wav in wavs] == [torch.Size([8]), torch.Size([12])] + return SimpleNamespace(audio_codes=codes, audio_codes_lengths=torch.tensor([3, 5])) + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache + ) + + async def resolve(ref): + return [0.0] * (8 if ref == "short" else 12), 24000, ref + + async def run(): + try: + short, long = await asyncio.gather( + encoder.encode("short", resolve_ref_audio=resolve, get_artifact_key=lambda key: key), + encoder.encode("long", resolve_ref_audio=resolve, get_artifact_key=lambda key: key), + ) + torch.testing.assert_close(short[0], codes[:, 0, :3].T) + torch.testing.assert_close(long[0], codes[:, 1, :5].T) + assert len(calls) == 1 + finally: + await encoder.aclose() + + asyncio.run(run()) + + +def test_failed_realtime_reference_encoding_does_not_populate_cache(): + cache = SpeakerEmbeddingCache() + resolve = AsyncMock(return_value=([0.0], 24000, "audio-key")) + + def fail(*args, **kwargs): + raise ValueError("invalid reference") + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=fail), variant="realtime", speaker_cache=cache + ) + + async def run(): + try: + with pytest.raises(ValueError, match="invalid reference"): + await encoder.encode("reference", resolve_ref_audio=resolve, get_artifact_key=lambda key: None) + assert cache.memory_bytes() == 0 + finally: + await encoder.aclose() + + asyncio.run(run()) diff --git a/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py new file mode 100644 index 000000000..f5d799dd7 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for MOSS audio tokenizer v2 projection modules.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerProjectedTransformer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_projected_transformer_keeps_learned_projections_when_dimensions_match() -> None: + module = MossAudioTokenizerProjectedTransformer( + input_dimension=8, + output_dimension=8, + d_model=8, + module_type="transformer", + num_heads=1, + num_layers=0, + positional_embedding="rope", + ) + + assert isinstance(module.input_proj, nn.Linear) + assert isinstance(module.output_proj, nn.Linear) + assert module.input_proj.weight.shape == (8, 8) + assert module.output_proj.weight.shape == (8, 8) + + x = torch.randn(2, 8, 5) + lengths = torch.tensor([5, 3], dtype=torch.long) + output, output_lengths = module(x, lengths) + + assert output.shape == x.shape + torch.testing.assert_close(output_lengths, lengths) + + +def test_projected_transformer_keeps_linear_projections_when_dimensions_differ() -> None: + module = MossAudioTokenizerProjectedTransformer( + input_dimension=8, + output_dimension=6, + d_model=10, + module_type="transformer", + num_heads=1, + num_layers=0, + positional_embedding="rope", + ) + + assert isinstance(module.input_proj, nn.Linear) + assert isinstance(module.output_proj, nn.Linear) diff --git a/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py new file mode 100644 index 000000000..916213318 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_realtime_talker_preserves_state_across_chunked_prefill() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSRealtimeTalkerForGeneration, + ) + + talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration) + talker.n_vq = 2 + talker.audio_pad_token = 1024 + talker.text_vocab_size = 16 + talker.text_pad_id = 2 + talker._batch_state = None + talker._batch_state_spans = None + + embedded_audio: list[torch.Tensor | None] = [] + + def build_input_embeds(input_ids, audio_codes): + embedded_audio.append(audio_codes.clone() if audio_codes is not None else None) + return torch.zeros((input_ids.shape[0], 1)) + + generated: list[None] = [] + + def generate_frame(*args, **kwargs): + generated.append(None) + return torch.ones((1, talker.n_vq), dtype=torch.long) + + talker._build_input_embeds = build_input_embeds + talker.local_transformer = SimpleNamespace(generate_frame=generate_frame) + talker.local_lm_heads = [] + + ref_codes = torch.arange(10, dtype=torch.long).reshape(5, 2) + info: dict[str, Any] = { + "codes": {"ref": ref_codes}, + "ids": {"all": [7, 8]}, + "max_new_frames": [3], + } + chunks = [(0, 2, False), (2, 2, False), (4, 1, True)] + + state: dict[str, Any] | None = None + for computed, length, eligible in chunks: + _, _, update = talker.preprocess( + torch.ones(length, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_num_computed_tokens=computed, + _omni_prompt_len=5, + **info, + ) + info.update(update) + if state is None: + state = info["audio_state"] + else: + assert info["audio_state"] is state + + output = talker.make_omni_output( + torch.ones((length, 1)), + model_intermediate_buffer=[info], + request_token_spans=[(0, length)], + request_sample_eligible=[eligible], + ) + if not eligible: + assert output.multimodal_outputs == {} + logits = talker.compute_logits(torch.ones((1, 1))) + assert logits.argmax().item() == (7 if eligible else talker.text_pad_id) + + assert state is not None + assert info["ref_offset"] == 5 + assert state["step"] == 1 + assert state["text_cursor"] == 1 + assert state["remaining_text"] == [7, 8] + assert len(generated) == 1 + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 2) + for actual, expected in zip(embedded_audio, ref_codes.split((2, 2, 1)), strict=True): + assert actual is not None + torch.testing.assert_close(actual, expected) diff --git a/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py new file mode 100644 index 000000000..1bbfaeeef --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_realtime_talker_stops_at_max_new_frames() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSRealtimeTalkerForGeneration, + ) + + talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration) + talker.n_vq = 4 + talker.audio_pad_token = 1024 + talker.text_vocab_size = 8 + talker.text_pad_id = 2 + talker.audio_eos_id = 3 + talker._batch_state = None + talker._build_input_embeds = lambda input_ids, audio_codes: torch.zeros((input_ids.shape[0], 1)) + + generated: list[None] = [] + + def generate_frame(*args, **kwargs): + generated.append(None) + return torch.ones((1, talker.n_vq), dtype=torch.long) + + talker.local_transformer = SimpleNamespace(generate_frame=generate_frame) + talker.local_lm_heads = [] + + _, _, info = talker.preprocess( + torch.tensor([1, 2]), + None, + max_new_frames=[2], + ) + assert info["audio_state"]["max_new_frames"] == 2 + + for _ in range(3): + output = talker.make_omni_output( + torch.ones((1, 1)), + runtime_additional_information=[info], + request_token_spans=[(0, 1)], + ) + + assert len(generated) == 2 + assert info["audio_state"]["step"] == 2 + assert info["audio_state"]["is_stopping"] is True + assert output.multimodal_outputs == {} + assert talker.compute_logits(torch.ones((1, 1))).argmax().item() == talker.audio_eos_id diff --git a/tests/model_executor/models/moss_tts/test_realtime_prompt.py b/tests/model_executor/models/moss_tts/test_realtime_prompt.py new file mode 100644 index 000000000..cb36f6dd8 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_prompt.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + del add_special_tokens + if text == "<|im_start|>assistant\n": + return [20, 21] + return list(range(100, 100 + len(text.split()))) + + +class _Processor: + audio_channel_pad = 1024 + audio_bos_token = 1025 + + def __init__(self) -> None: + self.reference_codes = None + + def make_ensemble(self, prompt_audio_tokens): + self.reference_codes = prompt_audio_tokens + return np.full((3, 17), self.audio_channel_pad, dtype=np.int64) + + +def test_build_realtime_prompt_aligns_reference_and_prefill_text() -> None: + processor = _Processor() + reference_codes = torch.arange(80, dtype=torch.int64).reshape(5, 16) + + params = build_realtime_prompt( + _Tokenizer(), + processor, + "one two three", + reference_codes, + ) + + assert np.array_equal(processor.reference_codes, reference_codes.numpy()) + assert params["prompt_token_ids"] == [1024, 1024, 1024, 20, 21, 100, 101, 102] + audio_grid = params["codes"]["ref"] + assert audio_grid.shape == (len(params["prompt_token_ids"]), 16) + assert audio_grid[-1, 0].item() == processor.audio_bos_token + assert "ids" not in params + + +def test_build_realtime_prompt_streams_text_after_first_twelve_tokens() -> None: + params = build_realtime_prompt( + _Tokenizer(), + _Processor(), + " ".join(f"word{i}" for i in range(15)), + torch.zeros((2, 16), dtype=torch.int64), + ) + + assert params["prompt_token_ids"][-12:] == list(range(100, 112)) + assert params["ids"] == {"all": [112, 113, 114]} + + +def test_build_realtime_prompt_rejects_wrong_processor_layout(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _Processor() + monkeypatch.setattr( + processor, + "make_ensemble", + lambda prompt_audio_tokens: np.zeros((3, 16), dtype=np.int64), + ) + + with pytest.raises(ValueError, match=r"expected \(sequence_length, 17\)"): + build_realtime_prompt( + _Tokenizer(), + processor, + "one", + torch.zeros((2, 16), dtype=torch.int64), + ) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index bed20efc2..555fec42d 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -6,8 +6,12 @@ Both variants share the same build/validate flow (``_build_moss_tts_params`` handles each); they are registered under distinct model-type names. """ +import asyncio +import threading from typing import TYPE_CHECKING, Any, cast +from transformers import AutoModel, AutoTokenizer +from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -18,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( apply_max_new_tokens, conditioning_cache_salt, ) +from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -30,6 +35,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): super().__init__(ctx) self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant() self._moss_processor_cache = None + self._moss_realtime_components_lock = threading.Lock() @property def engine_client(self): @@ -63,7 +69,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): # inside the model package; this layer only supplies the processor and # the process-wide speaker cache. encoder = build_reference_encoder( - self._get_moss_processor(), + self._get_moss_realtime_components()[2] if self._moss_variant == "realtime" else self._get_moss_processor(), variant=cast(str, self._moss_variant), speaker_cache=self._speaker_cache, ) @@ -157,6 +163,45 @@ class _MossTTSAdapterBase(ARTTSAdapter): self._moss_processor_cache = proc return proc + def _get_moss_realtime_components(self): + with self._moss_realtime_components_lock: + cached = getattr(self, "_moss_realtime_components", None) + if cached is not None: + return cached + + model_id = self.engine_client.model_config.model + processor_cls = get_class_from_dynamic_module( + "processing_mossttsrealtime.MossTTSRealtimeProcessor", + model_id, + ) + tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + processor = processor_cls(tokenizer=tokenizer) + hf_config = self.engine_client.model_config.hf_config + codec_path = str( + getattr( + hf_config, + "codec_model_name_or_path", + getattr( + hf_config, + "audio_tokenizer_name_or_path", + "OpenMOSS-Team/MOSS-Audio-Tokenizer", + ), + ) + ) + codec = ( + AutoModel.from_pretrained( + codec_path, + trust_remote_code=True, + ) + # Reference encoding runs in the API process. Keep the codec on CPU + # so it does not reserve accelerator memory outside the model workers. + .to("cpu") + .eval() + ) + cached = (tokenizer, processor, codec) + self._moss_realtime_components = cached + return cached + async def _build_moss_tts_params( self, request: "OpenAICreateSpeechRequest", @@ -196,24 +241,18 @@ class _MossTTSAdapterBase(ARTTSAdapter): params["ref_audio_cache_key"] = cache_key return params - # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ---- - # ``AutoProcessor.from_pretrained`` doesn't auto-discover - # ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the - # snapshot), and Realtime's prompt format diverges from MossTTSDelay - # (16-channel grid, separate per-step text feed). The - # ``prompt_audio_array`` shape lines up well enough with what the - # talker reads for short prompts; full Realtime support needs a - # separate processor.from_module path which we don't wire here. if v == "realtime": - params = { - "text": [request.input or ""], - "mode": ["voice_clone"], - } + tokenizer, processor, _ = await asyncio.to_thread(self._get_moss_realtime_components) + references, realtime_resolve_keys = await self._encode_moss_references( + request, + has_inline_ref_audio=has_inline_ref_audio, + two_speaker=False, + ) + params = build_realtime_prompt(tokenizer, processor, request.input or "", references[0]) if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) - params["prompt_audio_array"] = [[wav_list, sr]] - params["ref_audio_cache_key"] = cache_key + if 0 in realtime_resolve_keys: + params["ref_audio_cache_key"] = realtime_resolve_keys[0] return params # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 9d2d2d171..8d36086ef 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -44,10 +44,10 @@ def _maybe_prefix(prefix: str, name: str) -> str: def _iter_state_row_spans( - states: Sequence[dict[str, Any]], + states: Sequence[dict[str, Any] | None], spans: Sequence[tuple[int, int]] | None, num_rows: int, -) -> Iterable[tuple[dict[str, Any], int, int]]: +) -> Iterable[tuple[dict[str, Any] | None, int, int]]: """Yield state/logit-row ranges without assuming equal request row counts.""" if num_rows <= 0 or not states: return @@ -798,6 +798,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): have_multimodal_outputs: bool = True has_preprocess: bool = True has_postprocess: bool = True + # Intermediate prefill chunks reach make_omni_output but must not produce + # audio until the scheduler has completed the prompt. + requires_request_sample_eligibility: bool = True AUDIO_BOS = 1025 AUDIO_EOS = 1026 @@ -858,7 +861,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): # We still need a logits_processor so vLLM downstream stays happy with # something callable; a minimal pass-through suffices. self.logits_processor = LogitsProcessor(self.text_vocab_size) - self._batch_state: list[dict[str, Any]] | None = None + self._batch_state: list[dict[str, Any] | None] | None = None # Stacked weight cache built after load_weights(). # shape: (n_vq, audio_vocab_size, hidden_size) @@ -972,13 +975,14 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): span_len = int(input_ids.shape[0]) audio_state = info_dict.get("audio_state") is_first_call = not isinstance(audio_state, dict) + is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1)) - if span_len > 1 or is_first_call: + if is_prefill or is_first_call: # Prefill: read the full reference-audio code grid (channels x T) # from ``info_dict["codes"]["ref"]``; slice the chunk that aligns # with this prefill window. ref_codes = (info_dict.get("codes", {}) or {}).get("ref") - ref_offset = int(info_dict.get("ref_offset", 0)) + ref_offset = int(info_dict.get("_omni_num_computed_tokens", info_dict.get("ref_offset", 0))) chunk_audio = None if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: @@ -994,22 +998,24 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): if chunk_audio is not None else torch.full((self.n_vq,), self.audio_pad_token, dtype=torch.long, device=device) ) - # Capture the streaming-text list from the request (set by the - # realtime end2end builder; empty for the delay variants which - # don't take this code path anyway). - remaining = (info_dict.get("ids", {}) or {}).get("all") - if not isinstance(remaining, list): - remaining = [] info_update: dict[str, Any] = { - "audio_state": { + "audio_codes": {"current": current_codes}, + "ref_offset": ref_offset + span_len, + } + if is_first_call: + remaining = (info_dict.get("ids", {}) or {}).get("all") + if not isinstance(remaining, list): + remaining = [] + max_new_frames = info_dict.get("max_new_frames") + if isinstance(max_new_frames, (list, tuple)): + max_new_frames = max_new_frames[0] if max_new_frames else None + info_update["audio_state"] = { "is_stopping": False, "step": 0, "text_cursor": 0, "remaining_text": list(remaining), - }, - "audio_codes": {"current": current_codes}, - "ref_offset": ref_offset + span_len, - } + "max_new_frames": int(max_new_frames) if max_new_frames is not None else -1, + } return input_ids, embeds, info_update # Decode step: text_token from the just-sampled vLLM logit, plus the @@ -1046,13 +1052,21 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): info_dicts: list[dict[str, Any]] = ( kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] ) + sample_eligible = kwargs.get("request_sample_eligible") + eligible = [ + sample_eligible is None or i >= len(sample_eligible) or bool(sample_eligible[i]) + for i in range(len(info_dicts)) + ] # Defensive state init. for info in info_dicts: if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): info["audio_state"] = {"is_stopping": False, "step": 0} - self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] + self._batch_state = [ + (info["audio_state"] if isinstance(info, dict) and eligible[i] else None) + for i, info in enumerate(info_dicts) + ] self._batch_state_spans = kwargs.get("request_token_spans") # See delay talker: real per-request row spans from the runner are @@ -1074,6 +1088,8 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for i, info in enumerate(info_dicts): if not isinstance(info, dict): continue + if not eligible[i]: + continue row_start, row_end = spans[i] row_end = min(int(row_end), num_rows) if row_start >= row_end: @@ -1115,10 +1131,13 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): ) ch0 = int(new_codes[0].item()) + state["step"] = int(state.get("step", 0)) + 1 + max_new_frames = int(state.get("max_new_frames", -1)) + if max_new_frames > 0 and state["step"] >= max_new_frames: + state["is_stopping"] = True # Stop condition mirrors upstream: codebook 0 == eos_audio_id. if ch0 == self.AUDIO_EOS: state["is_stopping"] = True - state["step"] = int(state.get("step", 0)) + 1 info["audio_codes"] = { "current": new_codes, "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"), @@ -1127,7 +1146,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): if ch0 in (self.AUDIO_BOS, self.audio_pad_token): # Skip the bos / pad frames — they don't decode to real audio. - state["step"] = int(state.get("step", 0)) + 1 info["audio_codes"] = { "current": new_codes, "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"), @@ -1141,7 +1159,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): updated_acc = new_codes.unsqueeze(0) info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} - state["step"] = int(state.get("step", 0)) + 1 per_req_codes[i] = new_codes.unsqueeze(0) have_codes = True diff --git a/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py new file mode 100644 index 000000000..985f76205 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Prompt construction for MOSS-TTS-Realtime.""" + +from typing import Any + +import numpy as np +import torch + +REALTIME_AUDIO_CHANNELS = 16 +REALTIME_PREFILL_TEXT_TOKENS = 12 + + +def build_realtime_prompt( + tokenizer: Any, + processor: Any, + text: str, + reference_codes: torch.Tensor, +) -> dict[str, Any]: + """Build the mixed text/audio prompt consumed by the Realtime talker.""" + audio_tokens = reference_codes.detach().cpu().numpy() + system_grid = np.asarray(processor.make_ensemble(prompt_audio_tokens=audio_tokens), dtype=np.int64) + expected_width = REALTIME_AUDIO_CHANNELS + 1 + if system_grid.ndim != 2 or system_grid.shape[1] != expected_width: + raise ValueError( + f"MOSS-TTS-Realtime processor returned shape {system_grid.shape}; " + f"expected (sequence_length, {expected_width})" + ) + + assistant_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False) + assistant_grid = np.full( + (len(assistant_ids), expected_width), + processor.audio_channel_pad, + dtype=np.int64, + ) + assistant_grid[:, 0] = assistant_ids + + text_ids = tokenizer.encode(text, add_special_tokens=False) + if not text_ids: + raise ValueError("MOSS-TTS-Realtime input is empty after tokenization") + + prefill_len = min(len(text_ids), REALTIME_PREFILL_TEXT_TOKENS) + text_grid = np.full( + (prefill_len, expected_width), + processor.audio_channel_pad, + dtype=np.int64, + ) + text_grid[:, 0] = text_ids[:prefill_len] + text_grid[-1, 1] = processor.audio_bos_token + + prompt_grid = np.concatenate((system_grid, assistant_grid, text_grid), axis=0) + params: dict[str, Any] = { + "prompt_token_ids": prompt_grid[:, 0].tolist(), + "codes": {"ref": torch.from_numpy(prompt_grid[:, 1:].copy())}, + } + remaining_text_ids = text_ids[prefill_len:] + if remaining_text_ids: + params["ids"] = {"all": remaining_text_ids} + return params diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py index 36b2616e7..a828d59e0 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Reference-audio encoding + speaker cache for the MOSS-TTS-family talker. This lives in the model package (not the shared serving layer) so all @@ -360,6 +363,18 @@ class MossReferenceEncoder: await self._batcher.aclose() +class MossRealtimeReferenceEncoder(MossReferenceEncoder): + """Use the realtime codec with the shared reference cache and batcher.""" + + def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: + with torch.no_grad(): + encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq) + return [ + encoded.audio_codes[:, i, : int(length.item())].transpose(0, 1).contiguous() + for i, length in enumerate(encoded.audio_codes_lengths) + ] + + def build_reference_encoder( processor: Any, *, @@ -369,9 +384,17 @@ def build_reference_encoder( """Build the per-server encoder for a MOSS-TTS ``variant``. Derives the encode geometry (``n_vq`` and the working sample rate) from the - upstream processor's ``model_config`` so the variant knowledge stays in the - model package rather than in ``serving_speech.py``. + upstream processor's ``model_config``; Realtime takes its codec directly + and uses 16 codebooks at 24 kHz. Variant knowledge stays in the model package. """ + if variant == "realtime": + return MossRealtimeReferenceEncoder( + processor, + variant=variant, + n_vq=16, + sr_target=24000, + speaker_cache=speaker_cache, + ) n_vq = int(getattr(processor.model_config, "n_vq", 32)) # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate # regardless of its 48 kHz stereo *output* codec -- mirrors the offline From 5aadcbb0abae2e7065966483fcb90ac5ebcbd701 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:40:12 -0700 Subject: [PATCH 12/99] [Bugfix] Fix Breeze-TTS-2 repeated silence and audio EOS handling --- recipes/BreezeBlue/Breeze-TTS-2.md | 33 +++-- .../models/breeze_tts_2/test_instruction.py | 22 ++++ .../breeze_tts_2/test_repetition_penalty.py | 122 ++++++++++++++++++ .../models/breeze_tts_2/test_sampling_eos.py | 77 +++++++++++ .../models/breeze_tts_2/test_streaming.py | 120 +++++++++++++++-- vllm_omni/deploy/breeze_tts_2.yaml | 6 +- .../openai/tts_adapters/breeze_tts_2.py | 22 +++- .../modeling_breeze_tts_2_codec.py | 25 +++- .../modeling_breeze_tts_2_talker.py | 37 +++++- .../models/breeze_tts_2/pipeline.py | 6 + 10 files changed, 437 insertions(+), 33 deletions(-) create mode 100644 tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_sampling_eos.py diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md index 50e8858af..fb26ce91c 100644 --- a/recipes/BreezeBlue/Breeze-TTS-2.md +++ b/recipes/BreezeBlue/Breeze-TTS-2.md @@ -30,10 +30,20 @@ Four prompt modes are selected automatically from the request fields: | `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript | | `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery | -Current scope: greedy sampling with `cfg_scale=1.0`; CFG ≠ 1.0 and +Current scope: greedy sampling with `cfg_scale=1.0` and a default +`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and `negative_prompt` are rejected with explicit client errors until companion request support lands. +The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py) +also uses repetition penalty `1.1`, but its +[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py) +sample both the backbone and depth decoder with temperature `0.9` and top-k +`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design) +use CFG `4`. Those sampling and guidance settings are outside the current +port's scope, so the examples below do not establish parity with the official +voice-design outputs. + ## Supported model contract | Item | Value | @@ -42,9 +52,9 @@ request support lands. | Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) | | Languages | English and Chinese (model card) | | Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz | -| Output | 24 kHz mono; `wav` or `pcm`; streaming SSE `speech.audio.delta` | +| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) | | Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) | -| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0`; other guidance values and `negative_prompt` are rejected | +| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; other guidance values and `negative_prompt` are rejected | | Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` | ## References @@ -106,7 +116,7 @@ Voice design (instruction only, no reference audio): curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ - "input": "[laughs] Welcome to tonight's story time.", + "input": "(laugh) Welcome to our story time.", "instructions": "A warm young woman, clear voice, lively delivery.", "response_format": "wav" }' --output breeze_design.wav @@ -139,7 +149,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ }' --output breeze_direction.wav ``` -Streaming PCM (SSE `speech.audio.delta` events): +Streaming raw PCM (`stream_format="audio"`): ```bash curl -N -X POST http://localhost:8091/v1/audio/speech \ @@ -153,8 +163,9 @@ curl -N -X POST http://localhost:8091/v1/audio/speech \ }' --output breeze_stream.pcm ``` -Expected results: finite, non-silent 24 kHz mono WAV files; streaming requests -return incremental PCM deltas. On 1x L20 (greedy, single request after +Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example +returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta` +events containing base64-encoded audio. On 1x L20 (greedy, single request after warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. ## Notes @@ -166,6 +177,12 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. supported yet. - **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and `negative_prompt` return a client error. +- **Repetition penalty**: the default `1.1` matches upstream's penalty on + generated codebook-0 tokens and discourages repeated silence. Override it + with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another + finite positive value. Text and reference audio are excluded from its history. + The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set + the same serving override when comparing against those golden frames. - **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`; synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the 12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps. @@ -181,7 +198,7 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. | Feature | Status | Notes | | ------- | ------ | ----- | -| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks, SSE `speech.audio.delta` PCM | +| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` | | Voice cloning | ✓ | one reference clip plus transcript per request | | Voice design / voice direction | ✓ | `instructions` field, with or without a reference | | Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) | diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py index 3817ccb30..c7f2fef66 100644 --- a/tests/model_executor/models/breeze_tts_2/test_instruction.py +++ b/tests/model_executor/models/breeze_tts_2/test_instruction.py @@ -4,6 +4,7 @@ from types import SimpleNamespace import pytest +from vllm import SamplingParams from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( @@ -133,3 +134,24 @@ def test_non_greedy_sampling_overrides_are_rejected(): assert "greedy" in adapter.validate(_request({"top_p": 0.9})) assert "greedy" in adapter.validate(_request({"top_k": 50})) assert "temperature must be a number" in adapter.validate(_request({"temperature": "warm"})) + + +@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), "invalid"]) +def test_invalid_repetition_penalties_are_rejected(value): + assert "repetition_penalty" in _adapter().validate(_request({"repetition_penalty": value})) + + +@pytest.mark.parametrize("value", [0.9, 1.0, 1.1, "1.2"]) +@pytest.mark.parametrize("max_new_tokens", [None, 32]) +def test_repetition_penalty_override_reaches_standard_sampling_metadata(value, max_new_tokens): + adapter = _adapter() + request = _request({"repetition_penalty": value}, max_new_tokens=max_new_tokens) + original = [SamplingParams(max_tokens=128, repetition_penalty=1.1)] + + assert adapter.validate(request) is None + updated = adapter.apply_sampling_overrides(original, request) + + assert updated[0].repetition_penalty == float(value) + assert updated[0].max_tokens == (128 if max_new_tokens is None else max_new_tokens) + assert original[0].repetition_penalty == 1.1 + assert original[0].max_tokens == 128 diff --git a/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py new file mode 100644 index 000000000..0edaabc6f --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch +from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.metadata import SamplingMetadata + +from tests.model_executor.models.breeze_tts_2.test_talker_state_machine import _depth_codes, _talker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _model(first_score=11.0, second_score=10.5): + model = _talker() + model._async_chunk = True + logits = torch.full((1, 9), -100.0) + logits[0, 2:4] = torch.tensor([first_score, second_score]) + model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + return model + + +def _step(model, infos, penalties, *, no_penalties=False): + return model.make_omni_output( + torch.zeros(len(infos), 16), + model_intermediate_buffer=infos, + request_token_spans=[(i, i + 1) for i in range(len(infos))], + sampling_metadata=SamplingMetadata( + temperature=None, + all_greedy=True, + all_random=False, + top_p=None, + top_k=None, + generators={}, + max_num_logprobs=None, + no_penalties=no_penalties, + frequency_penalties=torch.zeros(len(infos)), + presence_penalties=torch.zeros(len(infos)), + repetition_penalties=torch.tensor(penalties), + # Text/reference conditioning must never be counted as generated + # codec history, even if the scheduler ids overlap the codec ids. + prompt_token_ids=torch.tensor([[2, 2]] * len(infos)), + output_token_ids=[list(info.get("breeze_generated_code0_ids", [])) for info in infos], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ), + ) + + +@pytest.mark.parametrize("scores", [(11.0, 10.5), (-1.0, -1.05)]) +def test_penalty_applies_once_per_generated_token_and_respects_logit_sign(scores): + model = _model(*scores) + info = {"prompt_ids": torch.tensor([2]), "input_values": torch.tensor([[2, 3, 4, 5]])} + chosen = [] + + for _ in range(4): + output = _step(model, [info], [1.1]) + chosen.append(output.multimodal_outputs["codes"]["audio"][0][0, 0].item()) + + assert chosen == [2, 3, 2, 2] + + +def test_request_history_and_penalty_are_independent(): + model = _model() + first: dict[str, Any] = {} + second: dict[str, Any] = {} + _step(model, [first], [1.1]) + output = _step(model, [first, second], [1.1, 1.1]) + assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2] + + # Reverse the batch order; a request's history follows its state, while + # each penalty follows the current sampling-metadata row. + output = _step(model, [second, first], [1.1, 1.0]) + assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2] + + +def test_disabled_penalties_do_not_read_uninitialized_sampling_buffer(): + model = _model() + info: dict[str, Any] = {} + for _ in range(2): + output = _step(model, [info], [float("nan")], no_penalties=True) + assert output.multimodal_outputs["codes"]["audio"][0][0, 0].item() == 2 + + +def test_scheduler_logits_preserve_selected_frame_under_further_sampling_transforms(): + model = _model() + info: dict[str, Any] = {} + _step(model, [info], [1.1]) + output = _step(model, [info], [1.1]) + selected = output.multimodal_outputs["codes"]["audio"][0][0, 0].item() + logits = model.compute_logits(torch.zeros(1, 16)) + + assert selected == 3 + assert logits.argmax(dim=-1).item() == selected + assert torch.isfinite(logits).sum().item() == 1 + assert logits[0, selected].item() == 0.0 + # The scheduler may apply repetition again and rescale temperature. A + # finite zero for the selected token preserves a valid point mass. + for temperature in (0.1, 0.9, 2.0): + probabilities = torch.softmax(logits / 1.1 / temperature, dim=-1) + assert torch.isfinite(probabilities).all() + assert probabilities[0, selected].item() == 1.0 + + +def test_repetition_penalty_can_select_eos_without_emitting_a_frame(): + model = _model() + info: dict[str, Any] = {} + _step(model, [info], [1.1]) + logits = torch.full((1, 9), -100.0) + logits[0, 2] = 11.0 + logits[0, 8] = 10.5 + model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone() + + output = _step(model, [info], [1.1]) + + assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert info["breeze_generated_frames"] == 1 + assert model.compute_logits(torch.zeros(1, 16)).argmax(dim=-1).item() == 8 diff --git a/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py b/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py new file mode 100644 index 000000000..8cd389d17 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest +import torch +import yaml +from vllm import SamplingParams +from vllm.v1.core.sched.utils import check_stop +from vllm.v1.request import Request, RequestStatus +from vllm.v1.sample.logits_processor import BatchUpdate, MinTokensLogitsProcessor + +from vllm_omni.config.stage_config import merge_sampling_constraints +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def talker_sampling_params(): + deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" + deploy = yaml.safe_load(deploy_path.read_text(encoding="utf-8")) + params = SamplingParams( + **merge_sampling_constraints( + deploy["stages"][0]["default_sampling_params"], + BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints, + ) + ) + # Mirror the input processor: Breeze's text tokenizer uses EOS 1, + # which is also an ordinary token in the generated codec vocabulary. + params.update_from_generation_config({"eos_token_id": [1]}, eos_token_id=1) + return params + + +@pytest.mark.parametrize("codec_token, should_stop", [(0, False), (1, False), (2047, False), (2051, True)]) +def test_scheduler_stops_only_on_codec_eos(talker_sampling_params, codec_token, should_stop): + request = Request( + request_id="breeze-eos-test", + prompt_token_ids=[7, 8, 9], + sampling_params=talker_sampling_params, + pooling_params=None, + ) + request.append_output_token_ids([42, codec_token]) + request.status = RequestStatus.RUNNING + + assert check_stop(request, max_model_len=4096) is should_stop + if should_stop: + assert request.status == RequestStatus.FINISHED_STOPPED + assert request.stop_reason == 2051 + else: + assert request.status == RequestStatus.RUNNING + + +@pytest.mark.parametrize("codec_token", [1, 2051]) +def test_first_step_keeps_codec_token_and_eos_logits(talker_sampling_params, codec_token, monkeypatch): + # vLLM's tensor helper uses a global flag rather than is_pin_memory, + # and pinned allocations require a CUDA device even for this CPU test. + monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False) + # ignore_eos still leaves text EOS in all_stop_token_ids. Exercise the + # actual processor so the first codec frame and immediate EOS stay valid. + processor = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False) + processor.update_state( + BatchUpdate(batch_size=1, removed=[], added=[(0, talker_sampling_params, None, [])], moved=[]) + ) + logits = torch.full((1, 2052), -float("inf")) + logits[0, codec_token] = 0.0 + + processed = processor.apply(logits.clone()) + + assert torch.equal(processed, logits) + assert processed.argmax(dim=-1).item() == codec_token + + +def test_pipeline_disallows_minimum_tokens_that_mask_selected_codec_ids(): + params = merge_sampling_constraints({"min_tokens": 10}, BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints) + assert params["min_tokens"] == 0 diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py index 045444ec3..a630b90bc 100644 --- a/tests/model_executor/models/breeze_tts_2/test_streaming.py +++ b/tests/model_executor/models/breeze_tts_2/test_streaming.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy from collections import defaultdict from types import SimpleNamespace @@ -20,7 +21,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _TransferManager: def __init__(self, chunk_frames: int = 4): self.connector = SimpleNamespace(config={"breeze_codec_chunk_frames": chunk_frames}) - self.code_prompt_token_ids = defaultdict(list) + self.code_prompt_token_ids: defaultdict[str, list[torch.Tensor]] = defaultdict(list) class _Request: @@ -120,7 +121,8 @@ def test_stateful_codec_rejects_partial_frame_chunk(): assert decoder.calls == [] -def test_stateful_codec_terminal_marker_pops_state_without_decoding(): +@pytest.mark.parametrize("token_counts", [(4, 0), (0, 4), (0, 4, 0, 8, 0)]) +def test_stateful_codec_terminal_marker_pops_state_without_decoding(token_counts): decoder = _Decoder() codec = object.__new__(BreezeTTS2MimiCodec) codec._async_chunk = True @@ -128,24 +130,114 @@ def test_stateful_codec_terminal_marker_pops_state_without_decoding(): codec._codebook_size = 8 codec._sample_rate = 24_000 codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) - codec._decoder_state_cache = {"scheduler-live": {}, "scheduler-done": {}} + request_ids = [f"scheduler-{index}" for index in range(len(token_counts))] + codec._decoder_state_cache = {request_id: {} for request_id in request_ids} output = codec.forward( - torch.tensor([0, 1, 2, 3], dtype=torch.long), + torch.arange(sum(token_counts), dtype=torch.long) % codec._codebook_size, runtime_additional_information=[ # Plain ``finished`` covers the fallback for direct in-process # callers that do not go through the connector receiver. - {"meta": {"request_id": "live", "finished": False}}, - {"meta": {"request_id": "done", "finished": True}}, + {"meta": {"finished": count == 0}} + for count in token_counts ], - seq_token_counts=[4, 0], - request_ids=["scheduler-live", "scheduler-done"], + seq_token_counts=token_counts, + request_ids=request_ids, ) - # The empty terminal marker decodes nothing and releases its state; the - # live request still decodes its own chunk. + # Empty terminal markers release their state without consuming an output + # slot belonging to another request, regardless of their batch position. assert len(decoder.calls) == 1 - assert decoder.calls[0][0][0] == 1 - assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [1920, 0] - assert list(codec._decoder_state_cache) == ["scheduler-live"] - assert codec._decoder_state_cache["scheduler-live"]["prefix_frames"] == 0 + assert decoder.calls[0][0][0] == sum(count > 0 for count in token_counts) + assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [ + count // codec._num_codebooks * 1920 for count in token_counts + ] + live_request_ids = [request_id for request_id, count in zip(request_ids, token_counts) if count] + assert list(codec._decoder_state_cache) == live_request_ids + assert all(codec._decoder_state_cache[request_id]["prefix_frames"] == 0 for request_id in live_request_ids) + + +@pytest.mark.parametrize("configured_chunk_frames", [None, 4]) +def test_stateful_codec_advances_decoder_cache_after_sliding_window(monkeypatch, tmp_path, configured_chunk_frames): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2DecoderConfig, + ) + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Decoder, + ) + + torch.manual_seed(0) + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(2,), + upsampling_ratios=(2,), + sliding_window=72, + ) + decoder = Qwen3TTSTokenizerV2Decoder(config).eval() + # Codebooks initialize to zeros before checkpoint loading; populate them + # so different codec IDs produce distinguishable cached frame histories. + with torch.no_grad(): + for name, parameter in decoder.quantizer.named_parameters(): + if name.endswith("embedding_sum"): + parameter.normal_() + tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder, get_output_sample_rate=lambda: 24_000)) + monkeypatch.setattr(Qwen3TTSTokenizer, "from_pretrained", lambda *args, **kwargs: tokenizer) + connector_extra = {} if configured_chunk_frames is None else {"breeze_codec_chunk_frames": configured_chunk_frames} + chunk_frames = configured_chunk_frames or 8 + codec = object.__new__(BreezeTTS2MimiCodec) + codec.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(stage_connector_config={"extra": connector_extra}), + device_config=SimpleNamespace(device="cpu"), + ) + codec._tokenizer_path = tmp_path / "audio_tokenizer" + codec._audio_tokenizer = None + codec._async_chunk = True + codec._num_codebooks = config.num_quantizers + codec._codebook_size = config.codebook_size + codec._decoder_state_cache = {} + codec.load_weights(iter(())) + + # At least two chunks after the decoder's rolling window expose a stale + # cache even when the first padded chunk still produces plausible audio. + total_frames = config.sliding_window + 2 * chunk_frames + tail_frames = 3 + codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, total_frames + tail_frames)) + for end in range(chunk_frames, total_frames + 1, chunk_frames): + flat = codes[0, :, end - chunk_frames : end].reshape(-1) + codec.forward( + flat, + runtime_additional_information=[{"meta": {"stream_finished": False}}], + seq_token_counts=[flat.numel()], + request_ids=["scheduler-live"], + ) + state = codec._decoder_state_cache["scheduler-live"] + assert state["suffix_frames"] == min(end - chunk_frames, config.sliding_window) + chunk_frames + with torch.inference_mode(): + expected_quantized = decoder.quantizer.decode(codes[..., max(0, end - config.sliding_window) : end]) + torch.testing.assert_close(state["suffix_quantized"], expected_quantized) + + # A short terminal chunk must emit only its real frames, then release + # state. Its waveform should agree with the unpadded single-request path. + with torch.inference_mode(): + expected_tail = decoder.chunked_decode(codes[..., -tail_frames:], caches=copy.deepcopy(state)) + flat_tail = codes[0, :, -tail_frames:].reshape(-1) + output = codec.forward( + flat_tail, + runtime_additional_information=[{"meta": {"stream_finished": True}}], + seq_token_counts=[flat_tail.numel()], + request_ids=["scheduler-live"], + ) + waveform = output.multimodal_outputs["model_outputs"][0] + assert waveform.numel() == tail_frames * decoder.total_upsample + torch.testing.assert_close(waveform, expected_tail.reshape(-1), atol=1e-5, rtol=1e-4) + assert codec._decoder_state_cache == {} diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml index 1d1b85391..fd75dd69e 100644 --- a/vllm_omni/deploy/breeze_tts_2.yaml +++ b/vllm_omni/deploy/breeze_tts_2.yaml @@ -32,8 +32,12 @@ stages: temperature: 0.0 top_p: 1.0 top_k: -1 + # Match upstream's penalty on generated codebook-0 history. + repetition_penalty: 1.1 max_tokens: 2048 - min_tokens: 1 + # min_tokens also masks the tokenizer's EOS 1 despite ignore_eos. + # Keep valid codec token 1 and an immediate backbone EOS available. + min_tokens: 0 - stage_id: 1 max_num_seqs: 8 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py index d2851e32c..03f0d3c4b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py @@ -4,6 +4,8 @@ """Breeze-TTS-2 serving adapter for synchronous full-payload generation.""" import asyncio +import copy +import math from typing import TYPE_CHECKING, Any import numpy as np @@ -49,6 +51,7 @@ class BreezeTTS2Adapter(ARTTSAdapter): return set() def _load_codec_frame_rate(self) -> float | None: + assert self.ctx.engine_client is not None config = self.ctx.engine_client.model_config.hf_config codec = getattr(config, "codec_config", None) if isinstance(codec, dict): @@ -99,6 +102,14 @@ class BreezeTTS2Adapter(ARTTSAdapter): greedy_error = _validate_greedy_sampling(extra_params) if greedy_error is not None: return greedy_error + repetition_penalty = extra_params.get("repetition_penalty") + if repetition_penalty is not None: + try: + repetition_penalty = float(repetition_penalty) + except (TypeError, ValueError): + return "Breeze-TTS-2 repetition_penalty must be finite and positive" + if not math.isfinite(repetition_penalty) or repetition_penalty <= 0: + return "Breeze-TTS-2 repetition_penalty must be finite and positive" if request.max_new_tokens is not None: if request.max_new_tokens < self.max_new_tokens_min: return f"max_new_tokens must be at least {self.max_new_tokens_min}" @@ -210,7 +221,16 @@ class BreezeTTS2Adapter(ARTTSAdapter): request_id: str | None = None, ) -> list: del prompt, request_id - return apply_max_new_tokens(sampling_params_list, request) + sampling_params_list = apply_max_new_tokens(sampling_params_list, request) + repetition_penalty = (request.extra_params or {}).get("repetition_penalty") + if repetition_penalty is not None: + # The generic speech handler copies extra_params to extra_args; + # promote this value so the worker receives sampling metadata. + # Only this scalar changes; preserve the caller's params without + # copying nested state or the unchanged codec-stage parameters. + sampling_params_list = [copy.copy(sampling_params_list[0]), *sampling_params_list[1:]] + sampling_params_list[0].repetition_penalty = float(repetition_penalty) + return sampling_params_list def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None: diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py index 7addb3934..860d64c44 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py @@ -204,6 +204,7 @@ class BreezeTTS2MimiCodec(nn.Module): f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}" ) valid: list[tuple[str, torch.Tensor, bool]] = [] + valid_indices: list[int] = [] for index, request_ids_tensor in enumerate(requests): runtime_info = ( runtime_additional_information[index] @@ -238,6 +239,7 @@ class BreezeTTS2MimiCodec(nn.Module): ) frames = int(flat.numel()) // self._num_codebooks valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished)) + valid_indices.append(index) waveforms = [empty.to(device=ids.device)] * len(requests) sample_rates = [sr] * len(requests) @@ -247,7 +249,6 @@ class BreezeTTS2MimiCodec(nn.Module): multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, ) states: list[dict[str, Any]] = [] - valid_indices: list[int] = [] max_frames = max(codes.shape[-1] for _, codes, _ in valid) batched_codes = valid[0][1].new_zeros((len(valid), self._num_codebooks, max_frames)) for row, (state_id, row_codes, _) in enumerate(valid): @@ -258,7 +259,6 @@ class BreezeTTS2MimiCodec(nn.Module): self._decoder_state_cache[state_id] = state state.setdefault("prefix_frames", 0) states.append(state) - valid_indices.append(row) decoded = decoder.batched_chunked_decode( batched_codes, @@ -353,6 +353,27 @@ class BreezeTTS2MimiCodec(nn.Module): device_map=str(device), ) self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate()) + if self._async_chunk: + connector_config = getattr(self.vllm_config.model_config, "stage_connector_config", None) + extra = ( + connector_config.get("extra", connector_config) + if isinstance(connector_config, Mapping) + else getattr(connector_config, "extra", None) + ) + value = extra.get("breeze_codec_chunk_frames") if isinstance(extra, Mapping) else None + try: + chunk_frames = int(value) if value is not None else 8 + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc + if chunk_frames <= 0: + raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}") + # Qwen treats chunks shorter than its configured window as + # terminal tails and does not advance their rolling cache. + # Match the producer's chunk size instead of its default 25. + decoder = self._audio_tokenizer.model.decoder + decoder._initial_codec_chunk_frames = chunk_frames + decoder._incremental_chunk_frames = chunk_frames + decoder._incremental_chunk_ramp = [] for _ in weights: pass tokenizer_module = getattr(self._audio_tokenizer, "model", None) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py index beaa390c8..765df5fbf 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py @@ -13,6 +13,7 @@ generation order without importing HuggingFace GenerationMixin. from __future__ import annotations import hashlib +import math import os from collections.abc import Iterable, Mapping, Sequence from pathlib import Path @@ -367,10 +368,12 @@ class BreezeTTS2TalkerForGeneration(nn.Module): return None logits = self._main_head_logits(hidden_states) - # ``make_omni_output`` runs immediately before this hook. It records - # requests whose previous hidden state selected EOS or reached the - # frame budget, so the scheduler receives a real stop token instead - # of depending on an arbitrary deploy-level max_tokens timeout. + # ``make_omni_output`` already selected codebook 0 and completed its + # frame. Return that same decision to the scheduler, including EOS. + # A finite zero survives the sampler's repetition/temperature passes + # without applying the repetition penalty twice or changing the frame. + # These logits describe deterministic codec selection, not model + # confidence; sampled-token logprobs are consequently zero. states = self._batch_state or [] if states: for state, row_start, row_end in _iter_request_rows( @@ -378,10 +381,11 @@ class BreezeTTS2TalkerForGeneration(nn.Module): getattr(self, "_batch_state_spans", None), int(logits.shape[0]), ): - if state.get("breeze_force_eos", False): + selected_code0 = state.get("breeze_selected_code0") + if selected_code0 is not None: row = logits[row_start:row_end] row.fill_(float("-inf")) - row[:, self.codebook_vocab_size] = 0.0 + row[:, int(selected_code0)] = 0.0 return logits def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -443,6 +447,13 @@ class BreezeTTS2TalkerForGeneration(nn.Module): self._batch_state = info_items self._batch_state_spans = spans per_request_codes: list[torch.Tensor] = [] + # The runner passes the same per-request sampling metadata to this + # hook and compute_logits, including offline SamplingParams overrides. + # When no_penalties is true vLLM does not initialize this GPU buffer. + sampling_metadata = kwargs.get("sampling_metadata") + repetition_penalties = None + if sampling_metadata is not None and not sampling_metadata.no_penalties: + repetition_penalties = sampling_metadata.repetition_penalties.tolist() for index, info in enumerate(info_items): if spans is not None and index < len(spans): @@ -458,11 +469,22 @@ class BreezeTTS2TalkerForGeneration(nn.Module): row_hidden = hidden[row_end - 1] if self._golden_dump_dir and bool(info.get("_omni_is_prefill", False)): self._dump_golden_prefill(info, row_hidden) - main_logits = self._main_head_logits(row_hidden.reshape(1, -1)) + main_logits = self._main_head_logits(row_hidden.reshape(1, -1)).float() + penalty = float(repetition_penalties[index]) if repetition_penalties is not None else 1.0 + if not math.isfinite(penalty) or penalty <= 0: + raise ValueError("Breeze repetition_penalty must be finite and positive") + # Match upstream: penalize each previously generated codebook-0 + # id once, excluding the text prompt and reference audio codes. + history = info.setdefault("breeze_generated_code0_ids", []) + if penalty != 1.0 and history: + history_ids = torch.tensor(history, device=main_logits.device, dtype=torch.long) + scores = main_logits[0, history_ids] + main_logits[0, history_ids] = torch.where(scores < 0, scores * penalty, scores / penalty) code0 = int(torch.argmax(main_logits, dim=-1).item()) generated_frames = int(info.get("breeze_generated_frames", 0) or 0) max_new_frames = int(info.get("breeze_max_new_frames", -1)) terminal = code0 == self.codebook_vocab_size or (max_new_frames > 0 and generated_frames >= max_new_frames) + info["breeze_selected_code0"] = self.codebook_vocab_size if terminal else code0 if terminal: info["breeze_force_eos"] = True existing = info.get("breeze_audio_codes") @@ -476,6 +498,7 @@ class BreezeTTS2TalkerForGeneration(nn.Module): raise RuntimeError(f"Breeze main head produced invalid codec id {code0}") frame = self._generate_depth_codes(row_hidden, code0) + history.append(code0) info["breeze_current_frame"] = frame[0].detach() info["breeze_generated_frames"] = generated_frames + 1 info["breeze_force_eos"] = False diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py index d1aaeb13e..2b538df01 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py @@ -28,6 +28,12 @@ BREEZE_TTS_2_PIPELINE = PipelineConfig( custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload", sampling_constraints={ "detokenize": False, + # Text EOS 1 is a valid codec token; only the backbone's + # explicit codec EOS below may terminate audio generation. + "ignore_eos": True, + # The talker has already selected the emitted frame's token; + # min_tokens must not mask that token in the sampler. + "min_tokens": 0, "stop_token_ids": [2051], }, ), From 63032c91c66bf7a4502744d138899554c26b5750 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Thu, 17 Sep 2026 04:43:43 +0800 Subject: [PATCH 13/99] [Frontend] Add ComfyUI Generate Music node with initial MiniMax Music 3 support (#7516) Signed-off-by: Weiming Liao --- apps/ComfyUI-vLLM-Omni/README.md | 12 + apps/ComfyUI-vLLM-Omni/__init__.py | 3 + .../comfyui_vllm_omni/nodes.py | 60 +++++ .../docs/images/comfyui-music-generation.jpg | Bin 0 -> 201933 bytes .../vLLM-Omni Music Generation.json | 206 ++++++++++++++++++ docs/features/comfyui.md | 1 + .../comfyui/test_comfyui_integration.py | 78 ++++++- 7 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg create mode 100644 apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index d0d002d0f..f9ee3b0c1 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -52,6 +52,7 @@ This extension offers the following nodes based on the output modalities (at **C - **FastH3 Deployment** for routing text-to-video requests to a MiniMax-H3 server with FastH3 fused at startup - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks - **TTS** and **TTS Voice Clone** for TTS tasks +- **Generate Music** for lyrics-and-description-to-music generation This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**) @@ -274,6 +275,17 @@ The node records which server the workflow targets; it does not start one, nor s > [!TIP] > There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node. +### Music generation (e.g., MiniMax Music 3) + +(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Music Generation**) + +

+ + + vLLM-Omni Music Generation + +

+ ### Chaining multiple model services (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**) diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py index 55b003ef1..8c1b03d87 100644 --- a/apps/ComfyUI-vLLM-Omni/__init__.py +++ b/apps/ComfyUI-vLLM-Omni/__init__.py @@ -19,6 +19,7 @@ from .comfyui_vllm_omni.nodes import ( VLLMOmniDiffusionSampling, VLLMOmniFastH3Deployment, VLLMOmniGenerateImage, + VLLMOmniGenerateMusic, VLLMOmniGenerateVideo, VLLMOmniMiniMaxH3Params, VLLMOmniQwenTTSParams, @@ -38,6 +39,7 @@ NODE_CLASS_MAPPINGS = { "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo, "VLLMOmniUnderstanding": VLLMOmniUnderstanding, "VLLMOmniTTS": VLLMOmniTTS, + "VLLMOmniGenerateMusic": VLLMOmniGenerateMusic, "VLLMOmniVoiceClone": VLLMOmniVoiceClone, "VLLMOmniVideoReferences": VLLMOmniVideoReferences, # === Params === @@ -58,6 +60,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VLLMOmniGenerateVideo": "Generate Video", "VLLMOmniUnderstanding": "Multimodality Understanding", "VLLMOmniTTS": "TTS (Text to Speech)", + "VLLMOmniGenerateMusic": "Generate Music", "VLLMOmniVoiceClone": "TTS Voice Cloning", "VLLMOmniVideoReferences": "Video References", # === Params === diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py index 47bdd9a2c..64020b797 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py @@ -503,6 +503,66 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase): return (audio,) +class VLLMOmniGenerateMusic(_VLLMOmniGenerateBase): + """Generate a song from lyrics and a musical description with MiniMax Music 3.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ("STRING", {"default": "MiniMaxAI/MiniMax-Music3"}), + "instructions": ( + "STRING", + { + "multiline": True, + "display_name": "caption", + "tooltip": "Music caption: describe genre, instruments, tempo and mood.", + }, + ), + "lyrics": ("STRING", {"multiline": True}), + "max_duration_seconds": ( + "FLOAT", + { + "default": 300.0, + "min": 1, + "max": 360, + "step": 0.01, + "tooltip": "Upper limit; rounded down to whole 25 Hz audio frames. May end earlier.", + }, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 2**53 - 1, "control_after_generate": True}), + "response_format": (["wav", "mp3", "flac", "opus"],), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "generate" + + async def generate( + self, + url: str, + model: str, + lyrics: str, + instructions: str, + response_format: AudioFormat, + max_duration_seconds: float, + seed: int = 0, + ) -> tuple[AudioInput]: + audio = await VLLMOmniClient(url.rstrip("/")).generate_speech( + model=model, + input=lyrics, + instructions=instructions, + voice="default", + speed=1.0, + response_format=response_format, + max_new_tokens=int(max_duration_seconds * 25), + seed=seed, + ) + return (audio,) + + class VLLMOmniVoiceClone(_VLLMOmniGenerateBase): @classmethod def INPUT_TYPES(cls): diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1c8130fbdcfef3b3d18add0a1b107ef7d7625620 GIT binary patch literal 201933 zcmeEv2UL^Ww(d`nF1>f@ic~?0fq+O80Rg2K6_6$(O@c&1QF=#d(tAgwDbjnBUKA8a z0@6ehX`zLeeeN53d+sp%opatjx4d_7WH^GDf6X=9H|P51;`HO`B0zgxLt6tNARqv= zaDRZ)IY8~VcmG#kzkSI6)d&3cpHJHWIx<2*qIe<#K7f#pfQXLZvv`y(}lMps03DT|-k#`?}Gs+r}oQ zX6Ckb_V*neot!6@C`y84E< z?;1NgySjUN`}zmQCnl$+XFh(Kg(Hwl%U@Sk*VfTLws&^-_Av*CzsW@a5dFSo|0dYK zlZy@~7a=h*5i#j+auE>v;9f*@#AnV)k}TYE{Oh3w7*F9w+R;X4@vgBV1LLp3s4ae;FL#12Y>;rOf~p;ClY(ydJ6Cq!x?;~ z?jnf>V}Y)&z}h?XLO{B=JlhV}1%0VgVC7kBI?wTwybulAE@Ej~{wO=vNB#zVL)nYF zY`|fFBM_zA_yFX)P2$PGnI$ByxfMbpDXG20kj_al5$^xVHaV){P;Z?74(az0574*; ze<9l=UxhVe+xQ0Eqf}};;()uzeJLu=0g^|TdNjyZ6bNg-_x$KOFeFr-crJ9Y{vKP8 zfM~DL+bS+;NH(Q(czslpop0Puc&M@$1iTRoEXK!ZsPLwTOZ_Cza(qPW=RUA6dB5ad znlveWVgiKgjfNH*-%s)c@ITXC>JE*f5eE@5(36Y>Xr&OkPcAx)dgV%8UVDvR8Dbl{ z;Q>u~$J4^f+SV={r}yX)ncZC~6}5AO@;$1dq|2O91JpDIaWuswVN0$P{$JHn33yo{ zCqrta(t+tz==dOoVHJSv)o|=etg_`>Wg}aToId{HegPyYKP!kci2T}|hP5=P@5!~! zEJYB{Zg0RNtzLd!T0IpzH^oa!{aNA$QJslNq47tb@7y6upy3Ci7LznCfIWFnmV^a5 zXC3=Euj^?qKJI7ZwF@BsS^R&JI)U(}RV-wCc`$hSvQ#5y5a?Tio}`CVaC;wbjK&wa z?{*4t^aP=?pSvVIMcCR#*jV|~mw>)`FNgj_Wy?IaP+>6|{{6RYZKML+8U$=@1R}2- z2YWPAX|z&2S2TmMZ&B3btV?!;(<~YfIf^dV=X=z3f_-PPo#0m<8?o; zARIU+U>BGmu+4*o=p3}40=>nANtuF|A8kT9u-gQA_xbm$(&H}#j~ww6#XOlx>6$C- z?9GyNkdo^Eac{{VbH|p5I8#-jy{)3?J*B)#ga*;j!HX_}lZ~A(d<|R*M473aQz`$_ z=zp5)7oz0S|@@D&X=v^~A6SGYei{s|Y*=$3y>9^!XV)G4bzBXiLozD>L04}jYdJ1rtn)|)3c}A9h2Q-p; zLxpF)Rd-y}eIia_E&c;whBt@)iA4C#`adb?{y(hJ7&^H`W^PH)TNKp|Ee;I=zB}Z- zMuAzX3?6H{^t}e z6@M0`3@(o)N5x8W$4jmZ$C947FLKgL_4N@m#GCur`rd<<5b!v4Br;n_Q6cG>kR_i@ zVp2i3Q~$O^`Y(O{KU0L3nnPZE>a0Uj9u1L1r^L>;cw8%`0e_34R@g9s_MQTn^w>u* z^m+l_5^VGg{=)75cmKpM|NpqA!k<=P|LPI_f8HY^PRYVPQ9lLzZBBvMG0l`yV7tR+ zmsPXPm(6OY(aK*b!0 zQOG$fCEY_1N0-VHztknIuw+LU|Cw&`w$pOhyy@!%>9W0yF2%TaVFZckRGB1CGG9Ih zw}MZ+ZBFRU8o$LoeU9en$tf_Vqe4F@Bd9{(qn}ESPzRCc3GA){0-0^}kP;9Eyi#kk zj|TUIw;)b|8B_2M)LRTN4spCDMRf5_7?PNck@jkUD)AESlL-Da+6jh>ci0Hn;(fVq zsPSiNsIqajbE#{S9)`g-Uce3s$9NX9u#HvNu#f=$OKLd4_z}hbg5j$;ffzZEjuapY zTp%4CyEyt1&|y=y6AouY2}LE&WCIC)fHorEZ5NGO!THu5d68G&B~HA9PXWzd71Auy z$ybUh#+QxQ&WG5Wku3*TF+YIN4UliM07=gR^xP8&+O~pMowcyXr@-zMZ0DMc z+#%)^*m$vU$jL~+S#$Lk!9}}+xJ&_~wlryd>i|OSZv?^-i6g;SKb(4D49AXtEYon5 zA(E(UrwR_JrTZ3SZy;xOmeeD1M=jzN zHeOC@jB-eR?A8cQM<%DhH8nL@RV?-$&Jc3iRp{`n@_&?%PKha-CJi$=@f&mBS5pkS zBpM&5#m}Er;B=cy4hzaT1?sa7FJA|e@IWHlAQlStJ7z-63UbSHifqSMTe3#gQn%CE&~PO>XiNE{sKp8`a&2e`%m^ezyeH-_JTR!nHxzkLp*RTL;=m)Q&uCw?T< zXJpIUj{DAv%NFgw1w5~H75PJU!8GY?5@K9NAXpW;s;?J7sQ)@|cKj5`dT|P% z$Cct@XPDz38TY~kO_Kv@FEojm2!T}@9dgjE%-^xzAB#*15LNpX z?xXua<`kYJNfi*)*Ay#sE`EfcH!5Z^CO7|&Jo8i1P+=@e;`hDRp*2u7>qEPK^(!xl z-Xm!U+xnB+slT>nCZ^tF|BtY_8Ek}kjofFO!wphwY)k;n?=i9%{;RlHK6w*X9EyFX zc?x8U0h<3orY1tDfM=D$zeByT-xpbsO)>f_*oQApl<@w9pnpyp#P{)m7(qN5jQ$bLUW>!%Yo)4{35@JI|ep3wb!>=zRg+n>YnD`c>%hK|m@z zKu=cScU1ESv;$(yaxMT#($GGA`tysrY{FG41Oj_NVz>5BI8R1*p9R~9!+m+k5=X+? zUx7XYJbJvK^%IT``)}%?r4xg6P1Sb=>`gYWgn6Iyq$LuTX9H2~DqZ;Y`<7L>DsNJys^EZ!N@&{kiIF93BiU=$?zv9fMUJFgN=aEd zkn#XSC8Yv`^m)@47XcE%?L4+b5(WLkwVF<~(XBqt^ngbY0eTX0Jd()J$z|}2a6{ic zf)@73sl{n#Bxv_gYe+&Ai18lD1dxMX>?UJzWEEV5x*&hgi0WFZv~3hxrp+TSn;PH_ z07zA&4DZ1X)>uvsU|6t0htQ9Eojs(<+yvp_M2>9&Hwfg{A6;5iyh(shcHo|kqExD+ zEe;h`z5>Z8Q6C5P_mPqYsq)_WM~;kvBtZ&=#1lLI)Iz2EQQVaA>D)9oAE^=(8Ty{z zkT}7?u2Uf9mdD?629-@=4t1V>lwe1EkoDN?mb+*SlcmTp(7&{8Y#XklJCv>kbWPA4 z*wQV9_A4mq=}BG+04newl;6iox&nyk55}VUgE@;%ygtEmR(+PJ^;G0R^yEau)glH% zSsUj#KS*uJ+(B7JEux+T2a)lz|0DfT;>_z$aD*cmu6KZZ$tiU{gzrYGR1YNWj)K7` zDL|tF5M{HL{4Fdi5gI0YQlm5HCLmp#RyL@ko}B>zr(lOC0?; z1&;mm@NogN3FTkZe(vuIpwg8*1(Kw=wV;WKlS@bw5Y@v@b^(|aD*fVd2Ce8P1= zlo&@%LmZ^i&icyhBV_w4)mI?}LW;R|bX9 z5=g(ZWJ?+xr$TV?I-{rXVV~xu8&5v4f>`+FLEPDkL1ZNPKa8erj!psaDZp90jE|E< z|3{cnq3fhixhK${$mtMXNc55mcm(VKfFeNzKR&-ahll?ufSZ(I-$UNxs-`!3Syy<%i*xhj-F$Ze?2L@U)c=!z| zdb|FZOI;HL12F;ez!L!A0=ntC3JSmK2>|z;LXWG%eyA;_XBG*wvb6!U_z(mBF!LvQ z24Bih#cRu|c!2M?k+jR%q9m-;8xz1C7 z72?qAPd)7hQbO_&V_s>XeV#*Ai=?|@Ihf?mw2Jzg9GA<~2=$mP-KmS4ZWU(J>e@L` zukVn7CvCcIFd{SG1kpr@y4WBrE%bT~h970K7ku&(yqE`OI|UqR7$B`33mjEYQMd#~ z#OIh1wx}8I&$ZD)hSmsNTnNb&M?bEHgM*EnGAU6piAXmBh+Gdl6}xyd!VDxQxErA* z?(fu*zCP_OJTP^?eAaTn$mQ|W3!g+cbr&BSsebaw(sliprIqWt!s1SI8hiDXt~>P& z;B<+QEC6JWHiA2Ba-e-WXOlx79LpU^Vp!T1XrWK-%jcbw4!>6AL2$~X3GgRhoo=D6 zuO6%8pghr=>$aiu(p57GHJhxZ)PjOtO25p`&pY2VcaobgKvnH=W2VrFds-(2;0Fst z{qQSW&s*OVh555P@-Xh4uRH4m^*@KnD;dmp#PXrOhr7Zznt6;+d*M&wPJs?1Pj7_S zVwd-IYbU>dQ~rRVfyaJwfrHylfjQ5`hs~$J8TykKySf&8B^V~uLi;=g zIx`rB+{0$ID>I;lJGN#U_MT$;{R$zN^{i^2%$zt*N9NC#-Vj;3L^dpPWMjiWU zg*ot3fLl|E{w}SAbkHd6x>)|)8_SCPx+&|U8%`all^2FyD?^ir0|(#Y zFG61umR!YlO?1}|JS(uBHXKm7BMJ-{OTko`Md<`10>O^1T49I%sm3d}>iV#}n8dCm zWjeG(MCGwdp%V`(ChP173PG7gf+=g?tblcxH-17V1(rHR!t1N+5^r()9!dKTEzpCV zPk{(>R4crxg>7tU;#dli6~?WI-ir0E!l+}}p(2>Oi@HLtKD{=K%0iH`Z#V}>#X$Hw z6tA82E?H2!TfAXSYV29%UKk!iR@KF1Y!AMSl z`>@&MJ%3cywYe?LT@Y6Mi!MUIK*E7HC&Bm*~`GR zRvkZ5=H`||#m0PXMHi+PVS+nN%(+t$x=Q|!PhPe@bo;oVgL#gYXyd+c@}h+e+ml5B zc8dK3L-u|vK#L+-+G82~C$rdOvxtPd-@B1JV{vX=j$uLU_ws}{WD1~8flgSMG9l_~ z9)jj9hP>$tLf5lnwkQl4qEBmO;nc$2!liiO(e->^S0=%s-K25n>*92gE}G|sau3pG zch<&Nf^PEFUZ0zDO-J;U14YeEShRE;UU~=xi>QTa>U9)}TqR)*yeeQy?f_|Hyo8{(K##8L4v(wZF*(UwD12$SobQ zpfY2Gh_aS^KMgYAi7&tf4O@5k-qfyk-N&74llBSDw~Cdmog?TUe~q#yWo?0IB0OA= zYW8AH9%}a-$}yK1ID;l`zNmGNLcmE!$cl!Rg=Q|Ej2@22*7+x;w&n%!ghG_vePh8l$@W zg&Rg6?ru>yC$lS89^SH+MlA$4nJ|_6>^0Ely3DhT*2y}4v~-&Oc6sJoRSbAB8%724 z^E@)eaCGxH;tdkWMDw@VN^H|lq}WRc2kGNFt*TdU>5WF;t_=9!XJ zj<}zfLPQan?b6@C#|T$zx^Fv8Cx&Tx49#$zlSXK@aY2omBf@ zg5tIWE`o7E8|@`w#?VA<)&khUX)pcFEiz0Ef~2_u12e()qM0-DU1n|eLX!Plz4y*K z=uFgI(U;U7j3=9O*@kKb9(7&y%y9@?Dmp)+|_IhkNyR;CzU z=g{z;IrFgc@Hp|`z@B=KP>Od5e?K+(U+PdrfG{O_IO_fl)@ATU;mBDUwN_pa4}$i^ zI4xGKY1i=pbdbPJPX=yiO5?q|xY+_v9B}77!{??5&IBEJEFS1aJq+cIy&{kp#Tj)6 z~hMG3{CHrK0lIeJl*CU@#ywD~*o-sx~gf zc`$}MysZQg1!*X-0YM%DKUX3F=bg}HW*ypb3}2sqvN!0el)wEb4@D%L_`wWcSxH#` zj|{&6R5~jfmxQQ1vUVpNLPB#E;uI8sMU<8R6@6EZrZlYlnWs!|`B-qf$ED!*b!sx+ z07v|LAc(Il`By3S&o9L&CVFyzbltxpH|v>dFB}%gdT%*-!HJENGol7xDORQ@cw`@A zExoXIMJqahr;)91Tl+ywihcXiN`V=E!Y=TWJ>#h{|7tJ-FP}4j3}0UXf3pW~vyWJJ zMRU71OWcy@e0B8b65Wbjf_rd$6gRahfVV^Ra~*XQ*)PaU|5 zr*~_Bx7qV2(&G~ZzbO;vy5>o+aoZ{J-5tr`ra45m$qU%V$Kj|)g!=zFQTX*a%D+w& zesP-eC!HSr-9G%2zkg+YB0nRU2mpiwlK-3M9mB?|MoZhL5+fCNxBWnqUphE=(vD~y z*Q@@WH2G^z2f&YSP8^7s-M%^HHK=n%VWly((vX$5vqfzwsEa`)npyVDlCQPQHL7pz zx5Mx;EuZXngt->JpMKaDp}m>?Cg?nYv18tQzoJxC;!XOx^efHK*NXQ8&I@;Zc<30I zcBI?z6!li!CAE)Y^;_jNIuZ8~?=Czlns;aq#1>4_R>(*0s z798cz1m#DY#6YFE3%v8R%U7k$EX~zRKlZxHetA}~#g`zOa;~X16_Ft8CiIouL~0<) z^O*AJ^06g)tJ~FKaXJg$nU7iM*k0&FQY!g>6;q}m7kc^Bt#IAb#jqr!_v2P&cD;)W znOoKUBo(oWLP~yNWzDE|^kqMB7Wn6flputbaSV5%+Dlu>kKg8_oo?NFsd(VeUU9>i zt2A1}diJF(Df(^xtZ%+Ug1yJ^cEoLup-hy<177Yt&rq)Xk-p^#;WslK%$ zELy-gr+qBO(~)7oA5C3Wln_>`6@er!4RG((6R~-&m|Od7y%9DH5B**#Pbfc;sklsj zNN1A=t21eHqs9(5;z|oBX%Qx5zY8nej7mNhm=~YAAHkuk^iY<|a9@9-+^a$neAz(=9vuL?x2-+cBe;^5}-`6;T1Q((5ypGm+FR2H>7D}BX_ zgW5kgT$yn_Ci|?3J}YD{rBYEKVx_^5F>uS~I;5&Y`8-PeoiDtK4dQ2t>S?`m(`t=< z`T6iG7{x%Ymi^^@gPtuncRI6*N?RH?PfCGNZ2`SX-&u6*O*9SWc3RQNnTz19sa}i` zKgOt3z&&Tm>xTi8Hl5|2mp7^9*Sh^)G7`YcP(`BfhY7&Sn#np31w%3oti+E@Wn4%I*5?`0MBV&-{^O&ME+rum+fYieU*4;vn<_QF8 zmti!!8*lf7Ba|psKRM6dp18XG>GOe;Nc8-?S$3R^VK1LH)wsIA!^)U}wo%d#3qQw~ zLQULrwUBD1YfHHo68PS>U&!WkpNQDMAVSn>f?I3FImODX?YDQ2%3^>RnU$D#SN)&e z#BEK8I&9MlOp$P&`*x4VCL34d*-)^7^|Cqud5}fnt$CBAtSAcPI%L8{ukJZlLZPx> zglmYDPk~p@<1n3nCAg!xB?3Tw#4@gjb2@rkfnKN^UNd18yhRkF=zbMv`Aqr3wZSJn^ygkv;x$=-iG(Ag-FzReRC*80Xve4P=*S_(W zm4zN(_+W$$O5)jHM;@P#$32B!essOjS0o}otWxrNF)61PqIm$ftSOxFi)!hW6&-(P z^b(|kxwn2bP6=l~KW})ITW$d8G8h-%G}*aK^abf|#WVGg-v*E|={H*`yw#n^ zr@T&qj82*_mXCT4{6j71Za9q!|HeaB7p zgJ1JM=G(uCqyJnnubAleL=X4nuc9&GeD5|VpY!_R7*=p&=d2yZ*EoWT2@l+Q;WzVu zQFvdAkx{J7yv|p`CmKaxO3=$CpM%%i{t|NjZm+?aZx{Hl9)2ShXU*|Fp|>acuiX|< zcMrO)B9AO$`hM#s=`*Lm&{u1%>O;azfwdBHll>_`YEBUh?0*7sdJNJ+yRl z|LNTN;l|TXiyu&-2EBu2W)7m$3h0JMU&a(Shp#xWA(F97)H$bsstfpNA9o~6XN^{u zD81Xd(hWX7!?1fp?)ac+A7<66{0(J~Q7hiGnypLvcJ^Zo^kJ|3fCJlFrn?I6+sid& zX0~8){;VsHt85Z0H^0dDT>0AkHA1;01joT(#91IC&#g|_NE*y%cW(Rp2(IJ(T&?&! z5Bl<(xsZ|XE5qL^KSNV`Gj1dAK|KFX2(yL-%H2&xlE%OaZzBrPy5_RJWW7#~A;}R= z_m=kLTP!4IfGDT_xod8@#f7imLTabjor=5;jH4f5)l4lx<_USaGT*brzDb|{Ucvu(o_x{D_z=Q@P)19(sIs{|2``#uoT?a z#xMT~{mSAKX1}Yx_+F4uX3saxmr>;6JMYQAbDnKuknKvmu^D5nA&{VJi*NkI#WP#* zZ@Bp<@BhqV0}Lm=JWuCU=O_3I9k>TzcMl)NTg0RXj@*cG9`d-n*6Mj~C?Jjs3>b$7a{!V&LezxI642p@c%V?HQ;jXS(;nuGt1re zS)d#;fo$BOu;{L*ox?TWE1i2`NcvJ{im%e`tj`}XM097sKUq~HEX&E0GMT+L|J z{pzIQ(E$mG3X0wp8&k8C`xb*o@`v2>akH=L zq^t;TsLns|cGxP9zd`l&%DU+1GMEmkDr~n|)M`d1gZ&h^IS!I+G=c;h)CK0(jk8!z z+Kzh#`|%=pHo2~weyi{tYM`d|RN;}sgQ;LY;{ee$9%hVAmp?IDs0z-*yk^2K9JV8x zXJ6)IqDwjA;JKWiE9d8XKS*UvhcjH@k~U*}W8)Q3BFoLcKb1XoLAX7i7ef`y1k`R` zou-VoMS!30&@f{>uT2OhoC3y*4iQ!26(G5@o8>nHNs1~o-MCJYvKW&Y>Yih>F6N#W zeMEIio7;MeB9vHI~?uJQwXu9dmF;7e;rxojj4&Af{a_XF46itEEiH(yWK zC4;Q2%$$<MbUepy0M;@17E>Phr-F)m2pE1`0Kjj(2Z!=2}FlRP!IaB~0!p(lrLac}H?39+g? z{&ct1g*#kte-|rcugqlTx$;QLa_Jz9@s2C+VvL6-*CibLI#c_!{sbyLE%0^qeKS=dhy*tqm<7W4&DI9^7B{27Osw| z@UV4C-Or*O%#!Q+pb;(Z7m{09XwA^@a%|kl$Xv~lFwg}d#G1AHG$=w5$ z-_iCd`1sjU{!GY?8G7AI=^j6%Kwqn)@Qe`15D2IIF?syQ;ITTO^V1;A-@rKx5cp~I z#JV4qSJ}EV4$CaEmcU0#BYkr$s1n-6D()eZ{ylynXY$bPey_EA)c&dn(f$#ZY~)A} z*6|wisrZ--EuOQ4kwZGXkfEv7sh`uqq!w7T?&&8tJqyLse-L{p!&V2w5v96HwQWyPdz1@C;SM>>4;erjSqoqKlS19o(jv$$tLgn4*(?s(I& zakvEj!9g2M9s1P0VB-mM(cSXZ!;HKGpC^w?3cjd6En2Y94qpCjZqlYpS~2 zza1)tYx~4N1xLNM7}v2fC|DOAGAm433GLS9#zsTA#?9J@!ps~S_zT9%ywA|=KD0Vc zs{S}>qo|Ve9=%B!H(t)9H&=KqZ*$A9fYiFI$utArM76S|ePzVFA_d7v zRvTr2l1i&$x_b;?5u@3KcBenBD*O)VmNO&gJoAdIou%3}d?p z5`5qdriXQ$_gvhhmmpndIkbQA+R_qpk8;k9K`ctKtTJI8D z-?yNk$u_3PxT}MCwWK5eF6$<1x4L5tmU{r1V5cCi&tV}t@S5%E^mq~j<@IMp-tW`1 zmlllY*8K}sVH8lFGR!tysVWw#j|8(wgq{Kp5@A?5lx6_-5>#>GKKg#F3%uFd^q3(R z6WQTUJGBK>MLrd>OYCY99c$w8M38UYAD4OMDm3|N)=IxnEkisIWqT7H+6iMZ#}IY4 z2qwUSZECHn5Gh%UJ=LfzcvZALO8EK1tz)rbM{AZzO4IJrwI8y}G4IMC55IGixor4R zVsO-j$`9!;527WakDb5~EIr$p)@}yK*u#6xHY|fx>*J6dH{NjGpiS@n_pYXvoy1O5 zMv>JQ-T5?PigPpytI_%hw=nM{H1m5jP0wLZ)1mrS=$!dv)9p%Y=TMc&?6>ukUQ$n$ zwm(9E%2JAgu4Uea5%1BA+Jss&;b;z3#wW?3yNKA(7Gl(V4ua$) z=OR=!8zb6bX3#{l1T8kAV4)7R7IQCO8NXG`c$wJii+JUCE?@UbP4p}r*5*vb}6Erbidub$o&)3Z{6IqG=y%P7a9b@WG}VyIq3Z-*{_=REB==xh;o zR?+9#!(B~OVRG@hKaQ#wetfCSxn_gt_B-5Nz_rb&g9KLzHMC$Smtm(Gy`%d=-78=I zZE!~xJ&sjKKXdyHNu$Yj!8t*Hbi{H?DssFzY?X6d<BehMbl@3Vp>UKXmJxFXcZCFnvp)?TCSpdbo?8x6 z=I!U_hI9pY&t4izI&5kDI$=6l$vf@kpnf@yM*buUdOI8A08j4T0JA(-CPgoF(LA%E zYI1ATS~P`TLK9uZRKedjt5({Pek;rUV*Sab;!#SH>;!pKC|`@vw7f#k%>~us4x?Zac#x*eI1K}3GuC+)nyBv!3qVpD&r0Wf=#g2Vld|+ zIx;*sz`uB%n>lgO=G#6xZ}0JYKeHw?S5%ETrU-rE>$p}W?^7ESOwD2}1%|zY{Ixi* z4wz)s!OTzFUw8h38L~eI=B5JYrY6efhAyHkj>R#ftB(CC-_b(Dw_9 zF_j%N*U{>@ftawm$iW4ArzG&*Bj${%FP7HY%GGxrPLd0sCXE;*e4Sr7QtTawR_=FO z-k^!JVIsw8FAkH81|Y@@dhf;-QKRVOu#iCNc;5$mCQoje%TH=#rnyd~CzxBBioLB% zo7Lt`5Q@;eO0cgP+r$Ikv{@9Yc3E5q<)%Qj6fL?vZy`jJMnSJ2W1k&vJ;PpDH6d}O zlhx;31+$7RC0|E>=DyH0c|UVTWx8^wT`15)xjnnP*vZZAM53F965|GcIkQQJYz(3eb56;*+CHaO#$Y+|{WG-igmehPFeepv|F%&C@o#d>$lp(RmI z>mrp@U6klUzBl$?9>mjYCu?JB(E6%qk}5Whqb054UekNj%&Y(>G|$+`l>L zO_^FHQC0P@z^U@#miTvToyfAX@Qme5gfF*Q!Q%u?hkLr$0g>E@|IpWgS}*K}L6h1J zZnk}=w`f}DkAZ7*e~`)oU6Wb@a(O&-kiY1-U1ocXP2=eyG79kqFFOm+{(YvgpmZD$ zRm=HjP5@Bx^8~H2?H{*4;BJ_Zb&=PWCWMZ>zgQ~&vX`Qg(I2q;;aet!Z?*vV<#c5p z-pM9>LrU3{vW!1nJ)-&Z93dW_4bfUG0P8C#cq%8+HJ+9Wp8iKTXr5OBP8Rc^$;Q2(cepcJC=uW@+8R~gE?8Q zKXDwjeZUsIBxlgP6i;ock!CVy=p?mVk|gpmG>I3cSUR=5yO$$3D@!?8wzgO5yPwQp z`Jq~a3EH2r;2gu`bB7uzttB{&UgQ$)0@H!iQ9K=)7f}n{ThG9(UR%$VImVP}kL@rH zC9^JAidK}wp6n!+eY(jg_o~?gFL>@3T~`F<0u3f}UOWH&J4Ng(m-Cq-hpOOF|*1_tLI@+zOZJ^lHAU2)78qoHh(s4{_D66g(>E}O&w}`);jgW zX6o2tt1ifFf#X$3b1gGs9$oPGU5v3pNi9uOD@&7bt}48W5^9JEhquzqewLl8jLe3N zbeStByFA)f>Pnb1b2^M}Zho)h$^Ut9sVB%HaW99BhE#4yVAj!Ekn$%i%Q1WzO^9#NJF!_MV)lX9t`U zG54`os$_@W?TL}Xx7sq~+QF0{iE;=`d%9ft9P~_RFE>9sOPd58<~>4bE4;3P$Y5&M zjTU2zXo;3;80ANW)Gj_MpGJoyc@ni@-GMI{H#Khv+wu zCW1@!##F79g!Q!77Y>`J(z-s{ucsu+cRo1lGB0m=Imx|qzp?IG{jLS{5a8&uR zMLx#QG-+1B>xige-Z8-{L(%>XZrtxprQHH2^KiyY)K;BKR=`4iOX5?Meg_XlZ7T`n z6+9~h1onl)NO;IU+c@o=0!-iXbU7?0S)3sqKRSkcbth-cQ*@t1m+3#+=eala?Ofz` zYga3izc6NC+=7PXSQEVxQSNY=#06(xtzGP2@;?-C)D+6p&nD$~QbiZl=lDWMj9*=ci{9akYqp%} zyX=$Xgv-|2Y#E?@sEQ{Lo0VmMaASi$>u^(o9YOy&h9c61-%CSGVk%JUrKZqGkNH`P21 zS=X@$=g*w~sy~bye=^5BM_AG5qUK8wPj!wM^iWlyP$r>Q_U__Y4&@ln4)&0|5{zDN zS5m$HnT9ZnPC5310Y?+|S~(MT*SA3&PIPx--Nc*_u?(h8?yHZL`DPx}n$`{4FsSFv zIv+l-Pi6_fc*oB|g0{D)^HEZGrP*x1U|nVT>&PPr7d!&af8UVc(BKrHB(VSoq?yDj z^Bt3*Xr9SbQ0~>Clww-puKe)kq7W|{CWwm;3Kp$Vm$!h?F$l4s+_X@ypY&F`EOvi= zUZ?(TLbpW0+tzWaDei(JB}Xs~RDBWl%-kg8m=wvwjLTH$5=>c#$-G!biwB>sGTP+k)wRiaR-sdx6vjDSWE?TXHP3#lQ&}H zOws`S0aO?9!OyaWh22XoX^#bLKe_e&&XI1uVoVNKb8qs>wZ`XrGbBHj1{2+pZfsxH zmAR(AA6p^(Dl9sf*^OVql?;(KQ}-vK5`~V2O|wF;;B%kEXIIo;k9|$h_9|1lY~k`$ z*^FC93+{7I1#BXKd5U`9y$`GLi8AT52BxXYj&x}NXHSgJqSpfHF+W0r^0On#`OWGK zH}$l+Ct~pIC4;)O@r5q)f)^|d<;kJkuy%L{gV$xy~(z7Mq+^u5qE9n{%Eu`cNV@Q76xv zpq115=7)G*j5d=!q2?Fdn4z0X>n9U-mc8kUlW?EX_R24gI8KsW)+ZCi<%;I1g8CZw zqJssySGOB#o15x_4HB5&hs=LpXZ>pWUC;zLn(`{b6d<>#AaV$nZ!5U!`*|j1%}LJj z`9+IX7KnaFGBBX_O1Q7?dZq&}o+~c)=(crx(jPK(8sU{{AOE(WmxtGLQOsI9D78)GXbQfW<@v1`3 zvRU`s-N~H{+nZi9VwDdH-^xiim85KJ>vzK#umYIxbGmQP4Y9{gh$AJ|8YklxA+&N` z<=oovV0A%Fq)~IDj&GV|wanOsfqPmAaocDnlWrKJvLNL9p#DUY8=ADIMRH;%S=uSd zS{SX-P-Sj<-5~Fhxy5?&jDe~`c$6~K`ON{x-G$`~*}w?{I%BWG6UDJNGw#+xj!}af zCohgLKrf+sL$^Q@i=FT*k(8HKuJCkyjbiIeq6welu@= zCIu}xveFe>>9*pzHQ7)V;a&9%cBT%-p(!FNV=dNq$_T(RFnbtKHk4v(TYt>C)$h1mmvWJJ!Kn^2RtwoVFDqt6e8LSsSDEz?JFc zxta-?k3wk(m!}fi)gf~?Cd8KCzD-)tudNT!{^}+Dnv_kQ@g~8ayhsPQubcu2(FmSG z+%R9ElJXmi-)~QxKqE*n(Fo3D9`P5-W?bN}m~Da-Q05F=lrhCS*Zj8BKe+2QkgGZ{`f@;#qjuNQUp}b z$dGI#oIEc)E{Vu1X{ab?vjfl1d|!J#{cfA&TmFLIKxEODJjdhR{$am+%~^0pxkR0+ zJbhx**c)zN+P2=GEv8l*kcvX7t_brNKz8sR`ujPw_lUV8sCw4IyhfUCuB^1KRWB4l zkyUIf0C)QyW!1rn{_z=}WB=IKRUfH5b|V=8R_i@i6EYARdXMX?1%hzNp$5ET^_AvQEhlc*?&5u>7@Af#+*B1T0d zN=Q_M(8m@OgeXWSgeU}1nnVaqDIz@)5JD1_9@Z)$E=h=I)_2FZckQZE=Y02!JMKN_ zj{UlZ2-CfOnBXZe2!3|%u>!|Dsdj%19RO!n)CZ7agYuj zp2AiJ0dfG)!MmOgM)f+_Ix`;p`&a+{hX1$TP2O37P?jVv;f}b1_Dx*EWLuz@_8#?9 zqB_Q?hLW`w%*CRi?9(kHRO=jk*6P%oh?9qGK_R>ZUMilf1qQzLc{hE=9aBiD*7WDXa1I6} zlF4NHiki#y$Jmyh75)Jj{0r;WtA;{>2z`ulsIe|FRTU1+x>2-LS)YY`7#4W6^TRRu z`|j>E8&T{wyX&X7<{%u8C5 zpS6dbImz$xl;NkJ&`)EuDeeLs+Rh*Yxmq5a=|7cqlj^scoR<)>RFCVJYw9XR{ycs$SNM#ke)zK5&D zo8LW5I#4TFd#FD!*Gcp9^ZoM<<+a8VEqsOd$@-SZ?fezHZ;J!-(3&T=#Fj?TAD(!A z653{e+-u3%t>u;&(x~Zys(1cHdv+hF-V4WkjX-N}?cI9q1Ny!a8$X7N!mf_R&rYTn z(0tu$54#G~{)i*-qk(($3wiOi^3|jo)>Y zadB`}99^O+AV+6Ogi+;Y8OKSbjn~s5jSpct!)`XA_dfY=LaaP-izs&k+&!YV;v~c; zpFVasi*aXMVgeCUW?FLR+h>y_je?aJ>gONqK%i7t0BCkR3+&$a+g|*zsmwn6zW3--(jg$ax9UxQ zgGN|{R&_i;k%{hz9g;0ebtWD$ogFm3bl0=Vl3ZY za2*G}BUrTS6BxcO+5KS2LP zDRc`uv|4%MkE=hyVpx>~Qk?dkEsfs13REsU6A#8Y4nPf?OSZk8eZg7fj#u3utlvB6 zEJK3UKxjfo7-V%j#Aa7oJPSjfP_mqOO+ocq7lX_iURJaeE%~LqSkJ9+NAS5bQM)<{ z1$H6n50)t3dC)7D@JYjsVJ{R%qX`Ujm!8oQX3Yz#c>w)5dg1BTj2HT`*L*;HsP-0m zM@hh6IaF#(iFI5>E_h2fM67}&KF3{dUyxMR+cms6=JDRt4yz#Ca9*6Ju4ehsU<~LqMExEv-mpA*>J-hnfe7`mT~Ic2u4&z`#m zZ&8#HeS@PSwXx!|`T>)n%hC)BcCdj=2@)M+_ zDiv1*dyhx3P!s&fqz5BPbRiV?rKdPNXV9jiq9TI7$!l{5$;EBlBN^W^qA}6y6a;;* z>+seogunFPP!Rq40vzcDgUT(XP)8A2j32(Y>UK;BP78lJY4`iDN6*H*y$OH*w(^eE zn4`+94}X7;#UmRmQvyAjs2FjIQGYddE$sNL%Zb4?#i!VO2z8X1s*ie)B(vRqtE%b} z?r)x`Ah@@kH+fAo+21<$EW$Rb!_==7V@Smz;SlF|OF0fU7Hss@yQXB(b11-@x(F?% z>AO0_Evs77I@oOUQDuF0w`gEI&yOBZ+%^%T){k#X-utSp@1_jzwRnB}B%ccb>Y_fU2$J^wmcnu1FuBx)S2<0Bb}mTwttm@K2T7JKo(#f5QwurjBLV zPn|OSR&#Rc-=h29&y~pK9jxzTWaYNqE?;p9DDhhuNLUa4Kcz~sc$B5O%IOt6x6Y~B z=IoUh=gy47kF{g_qywZH=|6j(!Xkwc_}MtR^3@=J%5YKz46GbUw%sFs9) zTN>?j6Pk9+CfL_BwD){^(RHlLbH}oKJzhO78bw^Kw|2kPEwa5nzwNLq!W7rU{IE{l z#H)DZ)#55{T=2=?t+_zgH71345b8_8ukAD|aM8}uyu11KJ|gsx3WkB6-m4`d4udSk zr=~g2*K-tW(DMO(_-Qy%xp2zor;4=&z3K$SQAQA!(#5Du`nSNnN`9cSR2lkEH~TsN z&(~-M2sK6(e`=}SnV%}Ru#bMKG&JZ<%2|`bZzFGtb{_#B3dF?e1*w}s!wkLer^?r_ zOk^K+F`(OiE5#^QOf)Lb`|>%T038XRb@sn@L%XMG{nJ4T`9`yYN9f{>i2F|!2a{Mn z#AwC+c%81)MupQSHd8=8r53v7r%G}XX81d9nA6{R!#_IrzjM3jzdRG2_!>|u(y{N= z4<8}^%n|N8!%{`!iE!ZEs#*jPBTeX;vr?-%#yXG;6` z3jE7{`S%J0{{QC;%#x|PO{dX)Sn3GI8@mK%bE1vti)rt%YoU&0D$NE-l*l+FD685- z^Qk0~w;TzR0P`79YQ@et);&}j%(%=_u}lb?PFfXh&}sI$mCZu$LAvF(icP`*OD}85 zW86YX2ku-5Kh9b`jbdCobxx@RITrp0wcYu$%}B1GYn3%`9Tf3WVrqDDd~d7wZ&pby zmE0J%OPu7R68~p&_&+$#vnFIB$#(%tjS85rZd&5H)RnqlUuO-nioP}E)`-WWs_OQF zw5|l&rdZtw0OnxI<-8prstE~i`B>Jxf7R5u+CFOGwp+_Z}Y#+~nzn_G*0wjXk~ACq7Y_ zgwLH`6!Lb*%&s@X>q0^|dE(W~$4}qvbyIyo`@mmT?kV4*xP%13{ck!-ucS+G4B=E^ zeq7$$P@_1BhFIKeF7(=xHDuA9%{1QJI$)`|Hq=`w!0ia^%7NzT8=WlKH0EY^^Y7 zqHVNzKf3m!=B((K+X8G}%RCtg+H)S-JZQ^qO3Ht1AM2lI>6oyFUzp<-J9J2PR~8uF zqjX(bxy+0y$t#Jwc-3Y7CztytOm5bH+jgWnr%AAX>%&)}-j?aNa@Xd08+RUI@67PP zcF^8k_!?XdjgP*S6KIWCbD_cvVoK&R#3^%|J{A?FKJDYHlUi3)v<6rjqzSKmUMS9XVX z`%Y_<|KlfSt2out=V`y5t6Ifp+!Ci8X8o9XS>-NJy%^mOW`%$wtlW+p;8H+~K*LAS zs66D}fA}9d9SMo;7r^0QA#q9-R3C1Nl}nmtG zgI7;x5L+7K`H%-FYq*K#TQ(vGE>QSKkk@SALQl#G>Jc%?-gY~5o)cW4m4l>=SzZ;l zeQsHP*Hi9dyUQhz&^A6gI)Wjao%{OgRm_i|R=V^ZS`OMd>8Eoy{Zl=Kw5{(#>-{q$ zZC7dU&6d=Ab~&5T+yZ=iVzUy(K5Oxv*}sK8^w^ZNUzUU?9y!|)S9EFu1T&RhVDH+t z)5U=Ze$LT*058iCf1NQ$W%f^%lNa`!+&&&Q=s^FEFD;*m-Gb~%`E?_9vBJ(|!O*jx zD!*z!nGLkk2>yzZe@eIfK`hN9)ExyH)$8beT0d2ciO)~a2aK3Xb;b{L8Xx+OL)NNJ z>oV2A#B?5V6;x{W`u|_Qv?01qio11V{j)HLBko&%d){pN$qPIB;%oVK{(lr#V?GCf zocvw!@lDmZua|R8UjAg9AG+zl))K zY)&fYJe0)G6{l{F_tacbuPj@Z$sPXX{>3G+o3lobtoDfIaokwkr^gPYJ$HRDK9c@U z%2ry(0&)cai!i?L5M96-Rr^edF8U1-ERL>PPwP2wG+KL#?T;0rB9ET4$=w^wf5|Xr zR5T3ZW==BIe2Gi32E(j{K`gBSBUFom_{<->Tv%=z`g~BiqJ?dipFkf~*KlwEkkQF> z&rcf;^?3O@$*-7NzPjg#>pIyc9Jult9`8HdfnGpGVX9Gcj#ZinI9=W1@fh%2GkuAM zla2EbBVuf+<_%PsB`H}HolwfGCN6afe15FAf%D0)sr)vPKwr(fL-sctD?-29%1)D; zb)<*VPdyv7!XB+?n_L>}20L%O^Q8FZlDnjyoyyM+TRaE50@ytxcLWy$KbvK++pCw7 zC5_=wnUhI$x$^~gG{3cc>4J{@tBna=%&3Xvf^?NIv-_+1$70=u=R%vOj?+vO1_jM4 z5O*;@?touluk!6Ry{bLRE!_xFVwIRa&}-Gg!vQ}`kF>)C`jNzwa4E&MX3Frh4W86t zWK@1$aIr{k_&bj9FU(by@H@{ z)vmC~p>XQKEI+cW+NStkh%av?uPF3v!2ZQoIZ1{`D}Su`uY~F?|41DGsDBx4IR#6& z<*KlIi8Kbn$Ty4e`hY=0{;i0&c8F4|44QJKg=Ef9XJ0U7&p)5+$ukqPf@Ss z=^!`6cTR(uiX!#;v`iwq|gNa|cs3i#i-1h}kG2AGSH}-!;(sOL<3-%O$x3^lRze zmUk#hLiSs;yoA)F?|ywBpK;9D&27%HWBcd70J0r`V0Fh5iwE3` zi_;x18(ZZZ_CFL3M#4H|Uy1WeT}%R87e5$pM=w)2+|ibhqXebvfEU~$*yxQDqg&@* zHI9W@9@T#*RM+vn8vSEks>#$KdWz{66#Y+uBsb%K3YVp2q+MV6N3n^&(H-~~R!pbA z6X)4J`)^qxjXI%=Gmlfk=6pEOik*m`1_@=R>|9X)l1AFi^8J6y3~683*LU%vL}MgQ zp0dT@KqfXs(Z3lRaw>D%aH@)Z#|x#g2%kYQa4h8B-NUL$`cPnbW%r8n797pYFssCq zKHcATO|hqYCciukWm~b*dOpN9b(V)b(0!d|9ZR*=W2B`nu?#gmRN7%@b3UttTTpiI zP!YGNV+9Yr?Kk`x$+g8S6>~;( zL^42Jwn`(|(CHP;8vTlu{rk!gFVM@Ix0F0eB91=ADK@aS$d{YnJ*lkep3X9n#Sf#^ z0pQJbcj?`_qpjqL+>hRSzHr+D%s1qpx@)rO9aE$F;;-yw(FZG!=g!QU?cT&!!+`)= zM`6@VwNE>W=!uze=@7Y^fOvST@nKmMXDEpoR!%fajXOaPPIE~^(gm=A%ieCml*93l zLl^uvc8flG_)q=vqQi_mE>V#YBG83s!IE-k`DTK&(Wr+KKQ48z{?f#(swb!od4v`< zMp&R?JvaiQ79p21R#zi)x7p|MAV1DbnK*qft0AYTy2OJTm$f>j_$3`WQ?(&jJ8i*w zm-TDkZas2k{$^(rl^J)P?HfXBn4<>LP+T>Awt}FUD0b%~s62lU(XEvUb#7Kk4|;Q8 z?PJKJi|8n0mrL6?X06^r^U)P9=AlNrTZ~IeuAVM;#(aUCm({PjyK%;|D_3SMpE+{$ z#@+p!DQ`^-8&oz&uYEG}fAt>^OnTi6rfMSA@VJCnJFY=4qcjbo;>^47^nqmhgu3RHW;a(gssP5H*&prz6y9^!oUem!PF8oL-mDaMLMBiW| z&61`dWr+l*(5B05;LN5$zT;#w$PH_wQIDuh&?`kNC_za)tq#yI9q}7YjbCo?JEju$k z2Didp+*pM(q!1jZ?**YgHo8K=(Gw!0nslT@42Btk-4T(jfhMmKEiAW`;EDR3l<$Xc zj@5ZX#bx5MJsU0utZy4ct!*<8Lggw!PF9Jh8lM2P^p_xVE#?M&PWeSCID4HXvMbY% zsY^45l*^Eei%+P&kbC9$GQ_dor5Gt$69+rEgl&gS1pJ68i|@QSaH60Fp@X_~72#Co zkjH}$x!HqV2d>Zhg6qbmt%!LYS$%Ty8XzK6GpL;z#Xn0JVJ<+30{zT5pg1Zf zkG$m!dL)EFS$zwpSO9U1Rz=9Q`fDDb!|iPE9d5-pti7s4Jw1=XtxV^He|E}2@;oqX&|-cF~ynSJc{UO;$a~sJyg>+V3G3} z_AE?LXb+3OSF+~S_{sOv+Lij?rt7PTPzCD^QQfJEvxG8S^}U!s3+xSm(7!4bMO|~u z%Ic-86iG~~=@H5;05NmJbNBBp?Bb&z{klFpmCbISlakj2w$~A?xd_bI(`?}WN>nF6 ztASxqx-SlFeKkMYRw$U6)q>Sh1l-%zRBj(CAuvXtD)u)`*2fvRw*nCs_YS+I(jIcH zQitD0HG492$qygocR~1ApI*63xoQ}z5wjCm1G+328J33X65N&tX2yXXGfH93@BJs6i!p&NABZjBMjHFrjsQFR!KomREq|| zY|uW-SJU?c=>BYsM6cTGTK7(Z+8+(~o#-6q}I~IRkVZMX0!O?hsZVxd@AD zJ!<2~ut{~V22tBDQMn>nK$}+rWlYBSYg2Hw-ph3lx&%#PsxgQOYJ?#XEiMs+nn<1z z3-#iufFNPcs_6py8BAD03*sQb)S8l)0TiE*w%C_H*KTp%r(Q(aftm|u9XzrkGa%Vh z>Faw$*exqtvAub92Rlfvv3dMplTsT-DRiU~W~6c%C9Kk&K9@ElD_U^|V)TI%qpvuv ztVgnfU}JAc^$pv#(uQ-LJG_>|GyR0jsxn*Mc%AEPcMKIq)p#=KZCd)x(#&a}eg^MR zNEXz)8kqH43SZbA%$i5AbS+T}s2ds8M0MObOcbuTk&%}-L0nvy)=NEtq)X8esaCKm zV=~!KLAn@v5t?c|0&?=ZwVHR=rPj~C*{1WP`xAMGZF_oNgC@!Mkq~XvmRhKsLc&I z)7RL^sZ|;y=epc$MZlh%^itS0(R~AqYXXt==os($J7!YkRXht&=Uc1 zqDW6jbUR-(?&d2WtY8cEmy8+028i#X+WJuVFDc@ zhS4>6T8i1o5v!NzHxn?Yn#=&peD;o{(dZMc(a5BU>H~Fq!Fmky^1y8B>X=$H2~%^h ziyucNJylq>*ttk^`r@cNa*mt!#x$WEVV{vZ@=b%HlDRn$zKW-{xL7@r+E_U;D1*8SCNL(uRRN$W2#4&nv^C{0uFa04gE)*FXO69tJ_6I$ zW%z_9=`37gcSIzx7ay}@bN8H@Qkv{rNx<|cMIZsZ=!k)Z2v5mSv-n9lFw zXoWivdMdvv99;?3 zOXl*~{@&&=^X4L4ffwmsjpl=(X9HKRT(dLW$A?@|F+3xd_&9wwR#!1^I@t~fairnn zqe}zC=rEc-%%vH@^qNqDpufs}iWSo;L04lU^`G$!>1tjikh(zb3`>G0#iJ3BRJBh; zPSnnQ7qT3cPt&G-5ET$cP5W>W{Gt@r$P`PBW=wsSgzutbhLXg}5_WxM;g=s$r8F+kbc{Ri_`) zPqU&$n!IJS&aA$JQ1A^daNT5CGg1VUdQL)sydbI&(W?cGi8G?5K!=MQ4!~*rNa{hh z;t2Hpmfb0nTCn~@siq~ZiencqX3(pa=5`H~blmRovcqYCes-lZcEg}@8)7VBMSySe zeHz1fix90)c#A7EUP)*0W_Ezu?Lp+ZTzWR?GAi~Gbyu_9DJfcy!T{K{T~OkKzR;Z? zq%~F@N;DLZ%1T4b>?}Gi_B_IgzYv#P1{yXBn-0(_>F%T7r0FTPAf7O~uQpC@IRKhy z;dO69$>PRZgOr`$CHV}gnO?P#eMfnK#96s~NH5`9Vw5F`X(-6SXZso%mcDr1DssNQ z5S0z00ZZivq84@=tx{w(qSO;Wjj|ww1G!4RkFJ()R~T|(RJ2`?n7-g6tuOHg=ifa1r8~0m_lVkwK$@2V06TdPCRdup-Y_TWNJ6_WdKXGEx9*NH?2$ zG+HPNWy&m}y3q(~u$alX%f8QB2Xo@tO5HBydyk;lRI~Dy38F!_;*8NsBvA7u_I2)# z%)ws6#2R;=5#`}PRKt)6IuIIIeNDACm&9VwbfLLVCHzkE_UYS7HH0W$Tv5lsE>Y+q zyKy!4v%tQq-)cc#Iwt;ib2uI`|^1DcziKUHdbjrnbpwvk4)?^s>7@syKi81;b0o#m(~bNGQyrxH>`O;dw=-$=R3VjMOh5 zTdkPj{DLsy!WyEsLL0Ju17JF0xjniMC7AXR*M;R@4smf59cts+2TH8tuqIwTND`^X z{XkK9`zc|0P;QrJV@Bsdhum`InTQbANv&JJmBz-`;psc&pw__B<*4OIwRK;X=Cy`+ae-o4$EiR>SHK z$>N4B^35GXN2>`T8-&8r;tblUP6MaKz3d#K^BJ9Y=L5K+7oDG+HxxG*AAW+0r>%HG zj^dfD0yx>*+kcz}=iw23pS{+7A5s3+52g$MRvNQbRK2ku)GuV4*MjabLFn)j?#&y|-WU>i$ks+4FA& zF#O;Dtu)F1X*$HJZL9@tm@sbB=IzM`f9FYPGyf+~!srF!g6}9*zp_ZPe*bOjGj@O# z4F9lEebX{J_tI5qMlg3qKYht~22raMIlpa{*A_vNRdm4UyN7_Y`70xXq(2QG%I(Jo z(l3}9m)j;u@vYqK<_2~}fz#)p300ZfZH9LbHG`If#qKm2qO zErmVk7CEP8^qIb4%7gru)pNPJG%p``Q6Ox>bz|W6c_&Vx9m}l!}OO_NPJgOJ*d`mIO2+eQr=)GbqV!?Up%q zTg;jH(`EaR3UgLK{(1tf)a_>1{-rS|`#d@nI%b3IS5W z{l*g{E)BS;{V8IQRahfMi&I~k+yG=cW-ZZHQp=REgJy2h&oS0v5O1xrBkpTzKQAy;<&|xG{mz;#?x;qUl$ybUXy{e z;`#iRqETZ37ERCs|96cMWw$A}@!7F@tGhSOIB&Q0!pO{CB;z68oqXzXh64~w^D>Zn zi5OS|&JB*Gq2(m!hSC`BxWVzA;>Sth#YMbFbBh4P%W|SrH_IrDcDhXN6*n55=6TF} z=+hU@NKD{!>Wp@trEjkZeCA7SSf3{WJyZ;!VyX&Hv53JrVI4VYVLFcT#Vs!1aW(50 zl3bItyJEy6<4bCsI3{H#={OGy7^^BMwg$jv#Z3 zcoJZw?5Y~8mt0Agx)8hr-0JpqSjf-ZeMdL+hX*bUM5TQ^>gnls8-M0X@i#zJ`g)qJ z@Cxq3i{>`;p(6_XlM*YZQGH;ruzH>uHNYhHG!h`rY#}k5XLTX@ ztw&ua*vMb4nKc2&88CaKKXz5wm{B6bcxP) z+(Oz4>OSZl(kS2=w~ds~li;Gd?zER5Aw3zqKq+n9*HLb5mF9kc3-vPt?9OCXX7ASc ze5Y~6YrXA8>V3?)jZ#!3AS?^DZ*I;nKPa`IRaxl$mLw}~@#&bXcUgqb*$`fnQ);!@ zYpHjem*2W-$XmZc+0@Fuv!ScbwWLk&y%J4zoK7j<1nAjp z5abAKy+=sh1X79_q*6xoS)?Y7)U^-ugc9S}Hd*0_o8zf4XwDyd`ntG~XLZenQZO%| z*VowQLFKJSWn_I?r$;MwS7J3^lR5)xtV;4?)IH>F!<#2Rr6j2Hqm0%;W9hl*dHy`C zJUE7v72Vq_I0AsbAdKyYiepQQ%h*Gkj@4WZT#)d@kS-v6^x22Job$D)j}z6q8CDJJ zhb6xQnxZXG*HcY}0prwzei?&_U7&8442r4ed8$Z2=JDYq`xKzvh|gsc4?7nb2*m1Y z69JA}E8HvTh5c3bfajyIPB1)by;G>%Ky^VnVS}nvn^L&sMmavjC#y$ce$S3=buMBWe;J|Eix9>JD6l0W<>Iuy(h)})$4CE>(FF^&YlLE!+&N!o)+V?~1-%|UTjrdL{(IwofOFY4|B%BE5dVoZA~W(SKG z`Jwfc#yk4Pqiqz<^Lq^ct&644DKKnlR={N6{7qHXo-if2OgwW-)wux42wF{;!( zf_yney~ZY^F5J@I==AGXZPay^?dg7pXgd|_kvkGX7)u9hpzwk&z5p5a2Dtwg)|hg1 zC?Qk6&&U{ZinREtqV3B#Sg_~~d4pqCrl&#KLAx~}^vx?aoDpya0Hy|VBgfj%Be6!L zR29TfZmqCpzQO4$jTNEgEnKfzl9H%Va78fxjJ^?sV&E3j#b>;zPQN=FS7KB78L^q? zf#x->6Y}Xrt^1*Ii(|a$dhUW~)V*+>+MZ6+F5$yIa|&)SZr3(4Tz*+NLz!UW%t=oEcxF zdOgT&#+brMYD!;1-3sqd7vq);y74w7kQWRs6b*DlQmrknl^^P!)`J7V5&~8k5!Wp) zhxVSfzNuUtvIT{?qO4kCr**A^QVT%10HW%RJx+EOvP$IiQumj&pQoNYjW7aW+rIDN zc4WBXI8&{gSCNh|1%dv&i#E@Ub_2od^%3k-uv4YCB`7@j;BKU9LUj??EtGoMGZ+(& z)$lc1-Eax;w~2K_m*fs#rZ6NeD5V zGilp{C86&s=PElzOX)tMTOuyGl0UN~!n}Nwd_MvRps+T4*y=%O9ul@a+Hgw^&9JSA zIYXr3(bBpdR*h5`OIdmp>uQ^fURJI%U0iW(tl7Av6z^GoK?6{7xQcNCjEZ9Rbd+5{ zpN7p~5IN?j$^xDx5>v}t!%I9;QUN~@-Qt5*ANU+)w~p;ixiU@^a+&I0fPh6E2_d{{ zR@hoxQ8&wN47AG1t`C0Rdref!>;!cU^es?}OXttS&P3={MDWQx1>reO;5OACjw>m5 z36c%aHc^cbb)e8ZI9(RPK91b&RxUCNsZCwO@A1j=r>PLdRx+Uge85DUP)(nZt*<0eRhK%=#RK1CQ zfX+K1Cn&8Uq^~UUmz(Npmg=F-Cts8DTWjgifw)t3Abd`&4kgOGrxTYWb$#8e8$u6a ztrw1ffTxZ(N$i_#G^^EasZ6fU?3zC;R^(=b5WY7`Ok5O0pVc*8;`q9Kb*mze? zX0bq0fs^GX9(29@j`CKePJI5f@@O zic{!YDUDI?)q3~*iBNLrQ?p)FBS!~m^?-@-fN*Nj7fDO`ZV5pbFf6CKXyuRm(i+4* zb@Y{Rb?Ay`rLTE~(I2M1BollpQpb9XBa7M!S3J6oJ8kC-<;vE=!diY;namhMkJRTV zw!%QHYnHNQAr&+ItRLaS+Fn3kn72~-R=F6aN7?PjRh)*D{2We|{gz6R?n;E5*ZpaS zxN#1$i=x?ow1Ba{^i`6Z={wLEsNBr=;cFnqHOh@JYqT!gHyEwPTS)srG{A19gzTSw z7&HoonksoZ)0xyw;xf^!LGMMS3`tE&BGc}Ok3^)8Y;kDMGfAwvW?`fcFEcW;nIeAl z=PfLRH~=qD>4G(891LS`6tgW?jjfu*6eqBZxN*9j-`^p)_qo@6x|p{fX>>roJn>IOBcWs; z`F5)Dbh2^*@(!EZ)YL?_UdPIkmGiV9c_nWpd^QS5lh+kPb5dBNSH#QA5(?g{l{-ARya-CC z5Gz@loPY^ZH3K*otMnAN?mF69dL>?IOFdTcf)Ar_DHk1c$x^nWHF5sg?x4SQD!_oc z6Cpxbl{nQ24gG_m(gx?`%EVNC$mmXJp5L#7`Nv5mHV+ySBe~O9lGoD;oHO&z8IC%=2~-P%6Q8YcqJP>EBl$fZ|q{k?Ebujb8eQ`pNdQB{GZfEfpa zj8?RngCI*Xrc0^} zaExo#f_gbQjHQP)r4pN;z*iXg{xYRWn{tc797?Q0Ekp3N#R)W%7PcjQwl~GLQg*Ow zFkz#gK36|ZnRRu|>kC=?!}gC_(2;e)qjRz-xnlQNOD}D4aBL7ET)A~R*Dfd+%;|4o z4_$MwQD=2;USdYuI)K)>oNagH%o#St$DJGhwcELMVhMFiVAh(aK^uhMg_}MSPR|%b z7cv#v778j5CMRn1YkC6rg-M7b-Ef#J@?I5o4q_Vw<0UaQ5KDHmLpdAYM`U1U6arXz z#Ce{{@DyLR9p;XbvXM)nCeGHfgMLZnrPScy9wigVvcr%lrW~DtptmD&eIfQ+96{I6 zk1zvx=qT8|3Kxf+4KpLa0?Y<2j2UoAEl~KQ-=H3osH^uFbqkUfC|7nu^#>8ZOiHRq zvYK1EkFLWhw`Y4E7l>_G3Z1~CLQeS{QoiyXy1;V$I(8{_V;SPwM~Fz^X&_n8{EW0( z&SO_Q!-qsQMmlh2G8}=5@zy~YB#HC*FR$Rl2<3`FccpR5q>DFNuTEk*x_|t@ysvh;+iEInbN5%KrB7@n3z{9m;%C~Ud!p40*w{T~FssMUQ;e;8!#=5vG ztE$1N?|axes@dw;YD$+|TTtl9@j7>#_8Ru|ZPmAEv{GN8jfB$sYj_<7^UbW3yBqRKYHO`E*EMFRrLLG_z-{?Ht$d z?vi|EDCZh74A<64?d!spH!0^+b$`bM%^1*2(gbcYEY(1`;gDj;qxUIi$fhx-oG7Jh zQ>36OUG!)?NVT{ixyEQy+sI|ZHC;kHG{$jkn^W+zIIefA;_X-0RlhlqX4VJUE!;C>Vwv7IQmQV*h|K?dH^B8! z%hk&=8>a@zwqg_oCNO*Mo_KgZnbScDu+mQmXIPpKGz)tAg}P z^g;VoatY7;F&mZOMEH?7A^1-k-Isq$qZ^=b(4?(FSf^Gpk^7;a%^DTPQ$E`knJgKN zFYhlpsb!slZZx94zCg{awRX7elycX(R?~4KW!vmOXkhPN!ctz_4dn5FyJI$QSAF;o zk%)r=W3`hrm#S&iZeFkzxAybd0~fBvUJu=FX~I%f$YPXA6Qu1zgIdJ$GqGCBehFK7 zQ|*A-^2aw(2e;4uUm=<@GFjFX@$k_pm0#7~$Nc)+$)(GF+pPBMnA(tOH&$~G6JX}z znUR;@kN)zpUL@UWvb3_UsY$i8{r~pVYjisUHk$G>ebZ%`^6c&|fCvh@lUuM({bB|s z$YLqCOdBgoGHEA|jdL~#M*X}8*2Qi+k zA+7H6ADZQ@dcbQ%o{v)=ARkSPVC%WT#p!VPNkr=uB&^~uI!9am5t(&>4;y6xvTv{F zi0U3F8@Xx{dG4;T<{(8N;H_z$#Gi#@yS(BeV%Fr97FD}%6yU#t0Qom%z9LWO&u{Cm zkI@}O{&-K&YX+wlq|WG|?xPt}Fk-^Ux13%QUsJi1(kKh! znb3@ZuFBTQ#>FH8J>O48TY9tKrFq1nGaoauL9a$>t8CQe&%)|^`Avcvvt)7b zxS+?Ox}!v&Yc<6W)nCcmrdsn741Y12jhsp1c|DYzpl0hXG zGc^9LQcq#tFiROt2SYI?%fYC!l8JaFsZQ&FUNRBc54sHv)6|gNXlsO?I2ruq7rDAp z@9d4AD!XxD080Bd{Ttos&llj?r!kWixcaHG-(G}*yjZ<(`%jfL7K{f8(XY54{PaIx zfoRn1)SyIvnfCf|6bp(0;hj(D%>*O^jJ$qED@4tIy#$BpoBof5K%e;Y{u%hoHH3BC zgWu?kA9Qz`8|uaXeC4EFYk%+s3Vtr%RP`efDziXMtpn&!KJ+*YOjl3g{RI zgW#veOjQ6Xd*5Z5g=uusfH0Nbb&g+DZks)5^M#T8lHPK;Veg00P4d@56I>y#P*1UF zTtDZTjd&VqZ{S=FH-wfc2+qF9_E&#=>b$i}e{eqog* z^~ESu)8l*&Hk(sdhz%N64sFRq@PdFR`+Jv95o92-yLpnj0b6I}AtCxnDZio>Q@$hgZ0X(O zb|>dy7q`&+9)3e&WZQ&%e!c$7V!Z4 zW~G%6;taCMn^RytEeCOWL71+3iS-SIwcL*BIg}9ViMCMUgH$yT^)*3 z&=Z6-psPA=BMx4yDvit@k54~L{?pFlR?+`t!~XiD69SpR^j+Lp&SDBf+#jK|R9Kb3 z;aQ_7`8CA>#1qULqK0f3F&3Wn2e?0oswI1BJPLKFGHO2DdWzq*p)*d zxW_I)8fDpt1ABnCu_-7{7ofVK{ywyz=Q8qKJy+hdSzfKzLCRSzs9k>hANE@M{m&lh zuP!aLP=fww(wE8U8)n9^v}s!@blEv17!t{R!L41T@T*;sS$)BAyjA|*n|{a5Vj@yk zw|{QSYziXBU2?IG;illLBEHommbyI^k1rUOX%8j!GQ`9MqQNF4xEklIm;tU6KwSf= z)`}P+Rk|N}0mW4EJ)F8t?oaDr&9@`xAXF+~fR3ex^wGmvFR@{;v~r6{l$||H*!vO? z(a^14rc>^Us$|$@T-fN2WCs0A2#h-3gvqx^`lsvoVe(h+_*VxD2s+fk;bF~DXiL$^ z+KlhD73+s<`I^{e2u{KQKKhj5PEqx>$zqliz1U0H=xsO6#u_T{%1$WZ<ustrc-E62!7JRAEPF);@3S$SYVu4m3n0^O z3HfJGk*PDF803`HCvt+YU}%v-1O5)`?UZBKZ6rD>2&Fa9f8-`wgMzNIQRtZkKB0V- zG&uxb_=iLN!v8Ng)G310U9m_|N7Muza_~eb+7iV+0M8`E5RJy{oM81n_Zu|yU>7ze zKuYX%tQWi(?p-kbbW)@RO24v$6svKGma?_DiDz)!ue7iR9Z~Lbp?AwGpu^PY`kmfc zubfvWrt2u{`NfN4vPK~R`ES1=*Vj8O4d@^ zul$!j5rE5Gfs(reOjY_7Ac0C-{Yc>|X+#{T))Hc8x`aQ91Ksu#kw!k%T`bZZoO+Fr z(iNDtpsTx6SrT;(r|NHb^Nf&Fda~0ZEoB28+V__3OPr?!ll$nUjSOR?zi+&&t~NmI zWh~|beV!pFxw$ujtpr04f~r9&6J=e)kHjrx)zn{oh8!a+bT2np@9x0hxRzH+7H;~7 zL4Eu`H>jh8xxk=q?StD`+Su6&TUaxUgJib;H}>8vs;R797Y0!f5u&sLLR3Ul#MsaX z#>5^GF)9iQLP}9uF=ZhVAtWl&3sEVd1SKK@B0^L;g^&OVL`6i}Xb{p_BAr=FAuJ)p ze|oWN@2_f~uj=eG{);olxxfn?Az4{-&H291`!tHD)3;RmgB(NlasZmwK|xadWtNkS zc#*bHWk=bZ#mk)uC#Ie~RjF%}`c_V^j^DZ6`Cvpl4z5z*yn zKu;UQ*@7}OgX@n2M&he4<_}L&pP??wPqn=MQej?!|Uuv9R7OHDy zZeh@r8hfw*sC%Z4@vVud$=u6=CyB?xe#&NfS9fsU*(0l&%Myg7Q1Bq#P--COp^eSN z^6`NmIHP%^BOsU1f@pe_t?;jCO>C4@o&u&#zYFLCa8}xw-!b-F%04)=T~Rf(87VAl^iFW@-^RaeqaEmC47}HMqQ!Cjwm$JiJ$8JLBQQqFIL*&f79` z-sVMD%yyU0QVIv+S&OB%&z9yK?5e1(-4yuo(8`IL^yCNqgO7=8g?G=NwhtboSNj{R zsun&7HOcaHJHLs0DD(#$U~UB90Lo7JceO=vH)RVM*1T~bg}2V>w&|A> zZCT5kkG*Kw_I&#DnlE|4P@|b)ApEtNMGL^8#b>M zb=&l7Kz;f?!f#J!RxF(napkI>76rX+^UNn}_fOr_k}q~%X4$!Ov-H3(1dXa)FEp0A z{iny(!C(ISaORZ+9O~Ps$+VIH-Tj6&r{82u5qgh((LR@P$-|nTh2#>6&_9_@oM(c?+<%Vc-LI6%5)q&Q|wnz5@sS; z>#K|Q=rjwvKQ^^%53vJ1f&ppK{b6jEPL>mlSMZUXV|U2jhbkZUREXm7n&Yijyqb@X zi~3xRx*vg&r4C_~pdC_`r5&FsYM75aBEEgUstRYTa7}iQD+^Y zH4t#l%30kbOnbs*G>W=W^@eH*kuLvK-K}&DJAaR|4o0Wa6R|d>q2=V9y|AIr9y7fQ zU;Ku(_IHDnw{S<@mBYt+$|C!mT1r7tHe0#3Yb>VER(d*3f;Oc@)PURz;zi=zQ*86l z$i@w){eI4HmK-6$AfK**c5< zLA3x_H+m|r*$i}i{|bSUBpYUEgozB;`edsnfeXC>Z=}SPcP&D0UC*)QB-AjkjHlTi zf)i5p9?$eQ?;amiI!NA)J&jzBbuAind9SUJtTwSDzimhKF^Q( zm^Yjt#lNUDIaes&@gX%S>RalR6Go};s-O_(@5tj$MGqYS$cFAk+SIJ~N2^B)m^&iS zqmJc1uKAy!cjznhnc$H)T|HzhtmbKT{lswV8nw%N0XmNPl-*#AbVbz%S8NYB@J?tCel?#?s;UUv^6p9 zj-}ZfWw;cr`OcS2;$1}yjy#>BmPeh8kiew;)ddgUPt#S}s#anJW@HqO+LZxfOo2oQ z=uAjy{cB+@I+MML+%v4W50jdK6qClYv_)cV>gDeMX=G6~^tl#dZo7Ea0@XasRN}ji zhdE=nl!fB3ousod<~#qhr-GyYeiQunnCL!+aY(N7b#q&El27`x`}S8(oh2CCDKgHV zS$%kwPeh++DB8gPrA2({rBK*kwEMf-iB541E89GTbBrVYl>@nyCaaAiZB zqO?YR|?Ke+X?JpYax9%=bzj;x1<$vWg=JgXN=O|Qt8M`BD zX+~&vW>!?t!}l)>{3rIeJNWiF+J4-4kYLxGW_e99+;}o`>X66QFGC?1o2Q|rn}2vL z9OoEN2uXe4)jrJor9eyBG8u&4SlmYajZc9OC+z&1kK1$e%kK+bIXefRN7prLs{ODB zpZ(fVwfK25koHOX0A|g*Bl!)3hUkNF*lW$?w$kIKdEj-Y$n#%ne>G9dynJErnd|EH z^0-gIGY9~)8@SB5g;qwVU=jR41#tpM1yj^A@5W7U$B)0nPpb$2178L;H?)YMs$r^< z66%=Z2_+@$yV^J(u$$k=OL$hv?jEmh?NZ})qP_*p|B(2?dDGs#Yt~;-Q?$h;>B~UE zc!b`~iF;9>CZ?Ur1DG3{W5x{CKUe;rxg+3wmTP>RCb)Im`xdW?OatA+?z;F7>AoI2 zJESi^IxmS>cprN2x-g~vdBH-XmycHdydwq!8X+O7nWm@JcqR`N2sNOraRxx4F~t>X z=ss0>MAcVD=6gVyS;fL(&WWCfJ9=)+;>6ikKZ+x$Ue%_A$c56$LF?*tLC1*FS9;rJ z_UI{!YnhKv6Mbcm=ulm{(C)$hJ?Dlu#LvlFKIhJ2?R{IXuN+X?c9QP#<*T^F^sH#=GV*T$F&b4R!G^mHL@Ck%E|q z{m**|bCd54nHk9w$|z0`2BR{K1t48EK-V{X>M*%1cH5z94y47k*e2;IhNifpxu}N0 zvm8TTvL?6D-%wBhdtJzmZN4oX##k(Sc1#)rg!eYQEZq8r%m#q@fb@IhMe7wapCNXn_=^H1`K|h#@szH^7PjUY6qnhl^5iW)v&fB%Ts- zN9ppWl9=%IQ|`(RmiY;)VF*2+Y~l@lMQF{@gcRk*2cLv~uL}wpBTPfz!K~b6*ps!lvrQA%ubK_w`iBkqRy zL{@$sJm+uLIVREx$=_4%Zrl~m?LRtXR+VR6zjliBm+zd`cwR=FG*fN?sCSr1wP`cj zw4E6E=ij>i<#p=tZ6ig{Uw-}P7yiFLx9R`7?@lTzA83{vsaiQ^fHDyUIWi=p-k^5S z!FnE(%29*;t17yfET{9uA*!GZ&GMl9A&|>*=6HX=`N!|V%+8KpeDb?<8vsoK-%?!6 zsTkm5W>a70&g!^>#FmmH6x+IR)w!rAkfQMoyLDHg5%Iq9vJ{8EB zJj!M`^CtANcWVZBEfU>AWR-Kw`tJ5Npne6TmVaHc|IHViFJ<hBkn%8AMD1u6a}&#AbhjLe}i$vKvVSxb@QEOzMi-SbY!2 z3D=5_@Uh!g+v+ZT-N#W3V-VY{(ZvC>afFV#M4esrlw(TKf_XkKBpge}8(_i-{372K z9h|-~=7+84=pT@)=Jzrag>|nVKk$jIo!U(^U1}CzSbFmAs$9#gdi+1wQpV(Zw1)lo z5QZM&t~d(qr`JdKCr()-zB1R#RCE*Q=T>GD3_>_&6rlHvpD<8vpVGQg@yVeNM9N%X zxy48r@lj@bEYSqLkTYAtin-F+CphlCUuX?`an|oCD;Sc}3E;Fab+9KA^Dd{6_C&85 z_TAsfUflv~?L)N15qS(56f#CN$b{TVro{#9>HqX)K4M?-3FZRdQa?GtUpT z1OPW*n>k|+NE~Rd1sr8G9U)-BVYrQ$gonbxXp4Ve;}bZdFFm8SuP4_wSg!72F9cAt zH4VIAUN~+MwT*4rqtuorK>U%9Ad~lJ%F<5Od?;XVCxuqn2J}nDF1QQ{+OR=9?LWaY zCcq~3Ds#A%moZG2YLClp5%6ZSmQCNs&MFXw_ycjh5|eb6l+Mtg1#@N-rZxi;T~}Ue zo4S&KprpRwJwIhL<7~3{!ef9CR$5jZDMVE|-Vou_(1 zHAdQjw*ed$9h6d-!4s028NiH063YWnCaytd=b{_64hQX~m6u$aURw7AVC= zh>=3EfL;#0{wQy%7E+?iY75D=2N?W_*21GlhN7UuEs`npT;HZMVG(kdM5Pf!aixoG zRmA(IG=cUzj3ITft87`DlKIwh|QE& zesUS{&gNXEIw>7!ag)tVnf6(^;oEJcNOc28Z*e=z77!&4C5Q-5+Q&qC!Ek7eKg;)! zOOc<9gRI2!7|E8}aNA`i5ke_`>4_w!veDBFr5#+t&&m#J_JWB`0O)D$UjhXsVYTu^ zJG7cpX2FF^}^CA3lZ zf+1mAZtb#iOgbQ9lbZbvIm5&swvp6FwcFddGt&yqMd?CV>u)SVMghX^X8gN_CW z+x;xEHVi7|UwU{}x2V}Td(Xaa{w^l_3hH0q zcD)31c^-K70}!7czAUp^f|`$TkCY3>@u@-vkEn|b?-5WeLLPt=POUgrn54~^9O&nH zq;s@#MTJgIxLMoFzuboZrUX>@d8_6;o7(HI+{r>11%slwto2f7_DZEDQc_fQ?GBRr zrE?cU5se&@;o~glK}QKvk|`1gxJ%|zlsfq<&xMo@#00CmD6Ue`ycX;)P106JhfM9K zYp2M4A^z*3s3fJq)C$#e%b8Gh0@fMc#&wZ5_E||1QygfTls%oQwOeLVod`+UZK=RH zJXK09HbjL5Q*^@fmD}FjLBhj=bQ3RmyH+s{`nhNRlbe$?eHR}hGElv#0#XU=MyE?S zOUiO21oCdoX2hwCDJu}~W706ZCPMHKuw3sWOkQ@JGv?7g6p|Uk_h>_#3VLj5OuFhBXz}tIe84gC z#)&5~xPFxOd=ce|3+)xB(DFYLH}SGM8M&-svsO$HcXSlcWE*rZ!{ecQqn<`j;GV zI$uiu&ktgnO!?G&q9NT=RdWh@Gq&-Imd&DtmfABfUON2i@{fe^tLcHijc2bRuip7H zdDE7Lne%kjj@-WkOm0!rztF#{ZT7+bI-aHauGR()`;{Axd(Pc{b6vm{*uMUk{5yJT zmKW5tHm}``KA;3$-?jhhcML`f+40c1Y|UK;q4@JR=`Qgve*Y~KpS)qAo^9P7JsSVV zWY51!W%DL4Yk&7c`&5gKJfH3ACpgyyTc3E~_x6mHY@Z=jbm;26eknLyeQ|Z!L+j3u zJ=Ua=UV#Uuk9Go&tU!+@PR$K%-vWh60}Ie`cUkL)L|V#7i?(F8Y3S)Mrl)Tx?P4Nk z_FJ%6R})y>N;FMU8a;5_6<4%$5wD*VuFs!2S=y5ya(A#MtZ~AnZk*d-#FBs9^w79? zMoVn${Z~KT7(Hd*5>`90%cf>tX=y<>&n!;1EIyQyfE2^V;_azs9>?zb9TgfLc)JUI z5I&c6ZpohODT;j;{E~Jo{h)CJwi(4s`$&pHBz4N`9ht_hp{24HD|aGAu(6CAwav4O z;#BCb^v>b6Jq&sC8?&Ny??-`0YLnCB!6$Vi2fk%0p7d#Vxz8;DSX)N3zXp&!R%}Gg zR${vW_eH^N5_)qnNt69)4+%0xVs=#8iLTR?so`_NQ1b@+7k05iS8awm^o_tOi}KN2{G{|Bh~AIT{A z6k0eViMrIRX_RwOwE&4|0!5>M+m`IV3a&~!Cfi7#kIsAmO1%axq}Z!Mw+-wlI2r;Yi9WDi$&AqGVS|A^(lZMw^D{3ipkKeAoWmhHrUz}{^CBb)|B4`%*7ee=_Q zI%M-Jy#E1zbL@}xD-+uOBmU+etLXplo6ssInS>Hi%5bsVGWgBm_rsvzBqS>R3VcJ) z9Q7!C9|{Iir4Bg@9A<}~K2=umICJRb^3T#LR>jlnn%9m!yR`L5(K`kNW7257Ho!vT zGT;RZ+eU}*;=4OBBLw5v9F@u8S|s%rl$~paQ7;yCd zo|5~Y_mmXMfetN$iW>3nj>{zS#;_16=<9F09_VQ;)mJF>)7RnAIMzF1Hf=^GO zt_ydU>}IAvWLZm>k1DODt>)fuq)41Zo@7Rgr?`fnJRe+iAnuTFTNc{_U~k)~is5_%cqiaDoT2Ixh|GAGX|75gyhEY{{||s03PNnjmp(fmJ0G=SO(xNWUMJ5WB^fW;fisgQDXA%=%&vj2Gjl)$*<`!BA!W-_`VD*8 zA@_FO^>{Tu%5CvMzt=BIuCIG#_{-)8=-zD1w5B~rxxfc}?iQG#O!WbFhu_r(V@s)} z$b@3cbS}%&W@HDXBl*2(gab`UO-3D{DTO7S=Wm{iMTt)Ra$E@7rmg))s(ntPxn4 z()Uk4y9BE8!We{^m5Y*PRAE9xhY8Y4e z<`7z;*G>zLfI1ICgk}q052(6{NGDb#Hniyx5r!nTAS-Fr82tiGon;I6@SDN0evjfV zupjntZT7OXA)@-pUR<8AfEj=K)XA>JFfm#3tM7L8uMEY<)bDCr(thwqj2Qsn8f`UE z#hDRRbM5fRc(nHwd|=xp8j1t{#V=s_P3;J4S9&3%;ywajWs0CBpv5hD$i0VnZ$acz zMqG5Iq7Jr{OEXi$!yG=}7IO^AC5>0NNmg|FcAeGrY|Jj;C&R3IYy-Z@-E;e6)sW5X6mwpz3W%ZIl#!iKqu6JP=o)9F=i|Yn4_pu#CSC zN%tf1qW1t+IBoX(p(@Yyey#f&aXo zF_Oi5a)fwEFx!IO?1eW}p1e!0mBbe*SJ~Mf>Zn|yvLI_UMW5r8ATKv<-Q>I5&v2`{ zl=v@8A75=xoSYj%L)KJ)DN{FobQ6>#7BtY4%orDN`pPgE)11oGM(9l{6wEG=qAqJl z9Lg9>#NFF(aT#{48K}V%RMt?><8?&b?FU3o{CHB+izH^EgSE-{S;@}GjV2Ml$yCIS?`n6Nv%qLZhoIWz`llPkZa~U`%5(hj zbS^#!twHU`F5;8oskZQlr3o?%Ch|Dt%w%+ZPn0*#0%_G`NG;5{HHUx>Gj$)r0Ja^Q zI~?e<`zhNyDhcsK)Y$^y8dDaD(I!{(GNW<%=TRPTi^2=pBqIt36bb3!G$0>W%~YbN zv+zMop;hxQWDUv30@yIb2|Z6WpZcmQ%5ouuxge*IB+r;6$`Lq7cBs79^#LN<+c}nM zUbpqj@IGUOo6ZHw+CT3G`G2}n;v@VMxl%tvLXXFXmJsGpJ2*zGu=^#YPKs=pQBML! zlAUae&UF8P;%o|B0io zi~$=Q?cF7rbSb{_bTL#yfH9-)U?-U1;*4?{6dcV8lqpxjoHIy!!sY}-L`s_lt1%c! zExbB=)rZb|?CLB}{T)TVceLvN?L#y@gWB?4?Gbat7H0Y)LBK#oR37`9QAWW}-=h7* z6r%aLj_tTr6n}S8+J0&~dj@na2Szne*SE{DCnHZvY~!Tz1x{($O%EvRpF(G^@e3>r z54#rl{AddYXiCFgh2QiIP7ppGFkSb6pTLqLq<0@MOeP7lcEIwcO*XvAfh2p5`KRaz zutIJ zr?bjCkIH9gE48HOaZZ%eNT)=ExUsdUt+d(WZD!roG=}U04aet z0@vG4p6*YPhu-w^aVtA0YihV5$G`4%xjvY9(4??F)OZ=RIlSrj;*7Sacqgy2|7Nb_eXsBhF&7R^Yj!g{r20vu=t9yt03| zVLoLStS(U6l!U-!>Q^q?rV`%Eh*_3XMmQ5IDe>6mnzO6sgGz_iun9MZ66wh#6BJKi z%bU9ls6Im4TirQ+cvKCy_D^5`UrU7&)=99<$=|Sar&c0#Swc13FTt%LdxGc4 zdde*ft+5#KT$Qa)!6FmzgBFQR+>$DUg*RDiCEn8s}mzdzw#hF)0n#C8>Qy2zz)KzpdM2GKWwwD3yn{Wai5jZ!Rm-Zs5 zTiVA^3x{gb^G}~pgK7U#^ZaXxIb|!wj2sa&jLE$c^hLG_B7NCIG6oeN)&V1=NwPAa z7$1nAid%SK69;(C$g85Hy&#HXkkH3RYFMN^#G{SSw_#7?By*I9X&r+Rjr&ebW|SQ~PBk^AVn zacQpUtfm{CR<3KO{Si+VmR+JwY3F{;^*Vavm8UB4<01epq0vl2(S&FiKKuKO|1;CgrQ-to`?2l0yfoNpa%Z!1qi%p zkjH%#&(EOI#fHBzmeijyRNk^QQy!FFQ48EvnHsY-qNfJ8sIaul&9fYKis^OD`8?>g%zY~ zX|B;qRXT(hxmhS~ioq~Y>5r8BnEz^pd4CD3u}II8j@p8-y0K{zw>(5;1z9FLmgJ|l zwcZrx_M9(0otG*6C_J+0X?3ns0@)^B=zJ`Dh0%$W!}q-0tRlT#y}b6-r>IS$|7LgQ z5Lv{jITS7S^hG5G)@WpCQ#3+g0uTL!KJo%Qjw&axlI2N45ntC=CIY4YR$_d$qvaZ8 z?E^%<4~db;FRE5SHI9s9@Djkm1WBM6bDy??L6d&luCnyhg7Z#S@Zv1mY}HJlsVjHl z;*&TPWm)jXM$UYENH`vBmH85pdI)3gGT;k?WZIFqHQ^6@$nIQ5MIeVMq&UIp<+KuH z=LFw>c}6`I2Q94GSx_8`*Di~6g;W>IE>Oe&?q&>HXT{?KX=bW7MCbrL5t#iriIdru zO6TbtDjYdS@+v#)LZy5rIZnc$Q`LoS`GekU{cUiqw5vK$su=%WpmZ%mQohLe>j!Jj zmE#jAi#q*-!zI;3$fLRw9}sJxxFt|+M69^}SQy_}3${H|U^H&tmXPhx&k-M1E%qm4 z_CD!wNZTf}=J0eqn^lY2&(?YKUG$O*lqMXgTai`NDiOey9Dujn-iOMd{rhU|GiHve ziSs1uBL1}eBK`@3$E5e;kvwMgs1Y;)IxnIYIY{iJ(2DnNHp@z7ndXJ|&w?g++Nw># z(W)YGNng}K7t#T2_*8aEM$ksAp(rsXjlG=Oo@)Uy(tz(u z?D~38>X1gwVGYS|NIqoi(Rmy;E`sD@JP&0`zKkSr8_^YYA2@pBqB(7ObpkF%|Hu%a z>78vQ@QkB(mPqBAK7%bWcV4PXY5%STlJ^@;Um{`mLt5_B1xt|ZOxoPlfD9uCg>a>X zn{vT~=OEgQr2!XhL9FvtuP1auq&MvZ-tg;68NiLGOWG-ZF=vYGEsfxU{$kAKK%%Z+ zsoZlalGd>CY}?yfuFu|L($r4MCgicymyqO5)WI!6Ox&J9#+wPLu4bb zhdLAT;Pd1lu!m{`jt-Q4sn>!^(UoD?U@zRpo5E`&`AwEuz<^R^gvhuZ%YcezHi%l% zB~;gwC9yEzeMRJrlegIG-Ym;F)P9z; z3yA02@?=R0zDo1R7 z=Uk?p#4l8}gQ>Dfj?{}O#-NdbdP{4N6hPEM!lXf5&a8f*yicQZOWZ}^yD)f2T@})A z4dHIX8ied&fJW&#^pY-HE(7d}F%xLClAl^p=icOUO6Lo6q!ro=IOnpUIO@EnNNLM_ z>96q8uOhNlgOEhWx&?(x+O&t+D{^Bny2Dg^B;rn+8hyu#;!x^hV2i56OQsz@E4P#Z z|5X2*k=AOVf4BDyv%wvn`4Cwz?lX|lxujIuLWA9m!!S6UN{%8HaHlI0P=v^c6`&am z4FLh)p~DT6)p1yJ7*xphJ&)R?epVigrrgDlsEL;3qaW-)i|z}h=L2GsyT?k|?_o@2 z+1yD~BgZVve+?oFxJxJRvK-dTiR@V8}yhz5;fW`cILH6Wkjht!rL z^dIJyr-_#sPlA#cLs{IVW^E~1QhH-ABjsFnT}2el2#d?MOKbX_eOnwGVmIEKMgSN? zj>vp4B6Ko=9weR)lc0M|>e1RbV`_)Ulx;yV>cIv)Kx`U_`q)@WaH%px;*QcGlB1`% z{f@A4Dkgrc=n0YxFQFNcY1_*mJ+ACfhCBrErL|J~ucyZe+QCtC0&dO7`uEH3_76YP zTHx@DFV{?kn}DBLDQO#$CdJ3&?2(&!N*m&(HKDnd&vLkv#L)&O62q-V8{#FW2Pcf_ z(oHCZ4pO~Ib)Ra{icVyXbF>ayXIASf@@a)rC?2qgmB#~BgU_;L&}&aMUJG@xOXTHn z(xSD16j}^6?}2)*gf>zKqZFC29mqdoY1~!4IRIJ;l|?yP7rO=OkRL=6Lm(&5YWV2n zCZ)@)kBKF3{a&ko%lq42qR`S7w0_%_SO|2~-XoqlP=w5e$DyMuP6_~{5aNPA(mk~#*k28zHEOf1UdTv*MX5zuduLeS(hkoAOf;4u z+XC{+i+y<8>OgmnYfq%bY>SnOUROcy!R(ejYTy10`1sryK?87ZWVxu^^@0N@*@D2sZSxyoWZ0I(m<=*u99vo5 zhU}FvJ3pWOEC>@H81Bpeq#(!U%1h+K1-bC&CEn>hmPW2Sh(_Zl5 z#0>4{83uqLv5O;b0RxE1^dk%IPrr(rxOyuC6_UT`%nE0li@i%vT)e*YD_TW%|JDFt zqx&;5bFCeJ9wj%bzN+U3T1_x-U0WH~Hf5!e9^L&LueXrPXRHZYyY|JG;~%KVPXQV; zbGJA*)Rz|3>TKAYvQy*Lw_kbc)5BwZOR6oLHF9q*?EdXTcqjGawwvmU7Lcc3ZPQwJ z@zz}U+9l`Lo7HCLyrvJ(aYvuxJOj-0o?E?su~epIwE(;DlJ@rAo&M{v&!e|<87^4# zIRC=WyZBa%=Iz~lg5n1+;!gSGz3)*Y(%CS_H5jcFyC?TsP0^%83O(Q};iAbw@mIsy2BcSYwEg5R({J(Z@6XTDkDX0dS9ajROm2 z{g#dwI%DCETA`S2dYxiXc)h?aXa6_f5G&O(lTGK!F`LLX=})oVF`F%8;pmU&vHOb< zFYP&llX^Gu_+cbl>k71RQk8mdsAhDGzBXlvz$N~ruq?JXT{NH z!=Q+AK%10ulE;>x?1~K2UBcfB8(&j9g85MS=3xe*v$AewyV}#binwj%Pbz!t#O@^} z+GoF2Ci_k~j(56$1U~9P$PFl0{sK~zM;KlCB_m35`} z<$mNvL7#|%eG0w4*8*{x%n0!gD4kH36a9AGv*O453uxz<_qkEDkW*DU*7u@ID@UAG zNGwf}PAPu2(yNDU8)D{vG4j0Bma(AhnIx9_I-J;8^D$o5;zuCBPVwUpP)`FXUz#GGAXmSt z24<$MHg95RvMk`h-aSYRg9(5Tc z+>`b5kq{SL!P$(N;~oYIPrx6&{E?)E(Dyk-lNACcela zDC!_|NEFAmKvv+;x1IQtr1;XYwsTSF{q({d~sFhVWb)=Ym=ao?edaY z`@R(3JSizz&aBvA7WJ%BVCJ4P-0DtPQMgF!7>rQ_AwOX%&C!=R^RTY>Lxrf31~G&F zBp9~>iY&|k4gD$;+yaQ0#M-mUVL2tNq%L0BZzOPFB&`chE*e@B+#ceTYZiB)OBNsb z@>{hG!_|OjO+C3#&++paANi2@J7FQ2a|DT# z;f)YE@4&-J`4yJyLxkadeND^!uE^G{vIDh;=ij{IHR#a7C(SKY!aBqK zn`iyQAM{gy?dh~1x`YLk-P1X=kY$kW81b=J^^_qpptysKi%U?(;oxfs0Oc=?J zf_83XrZOtLM^I9@ntUp<>pnm2x9fqLz6|9fkTN^>9X7R$hJHe@LgjxkVOA>&#l?M(_vexYp2d`^-U@q zi#1Wrv4i7wAe#y)Yb<9Vm|}#(<;-LerW_GhMN+rYMyeyC1E+W+nNFJbRG6gyqSQb- z2F2FhipJHJkqO%bWmw%vX`3$D{b(;8OL#Mja4fCs~Q zGyx_zrbr3Te-$v9Wj2Z`Y*Min5_~ewB4;XKkH-vwgwa?>BD7=oi^KNTlS9a|U=hC8&TLC1sZbP(Y}ak~>70NiJ&0 zUENW(x^l`q@M>L?{9>eeD#qfIE@?qB<}=gYl;vch%qMiR+)H+AI8KZJG1)A40MVU5 zzsmv3wVX8a^*$0)zZQ&-&(*6uF00pkJkr0ppQ4F zZ-zy#I&M>!88+PKdw{ccBrQKh?$yP-_ax<{lcSLI_SEO4nZAM79sv3+ejb+rJVM?+ z&sSNa&+CBxb++o}>+Q+S8v~fmqlYhh47h&fGgco(uXvpH=EMAf6W7ngUAtnbx5b<` z?=2IQeo`mex0wrzw|rMyr@G01Jrp4-Tr0wQExyM3)RwV~@@(gNy}dq`uJSV%FX-*U zUYW7}Uka4F^tJyf6;&e%OF08l)YLQ@*s|>BMx0)pWS+ETe)H^g&9fJ8UUcc*Ov}4P zUycL92q-s2Iw~W>`P+|O{!O;D$<*ffm))1#WX0HQwNp`!lEP)ud2=sZX+c}PNIF4z zf92A;jh9?f6!N%V_}2U+>@)p6+X{mMN@N(?l~G_H+8iwwB=Ze@2d05Z%^gZ={g>k# z20NEBif!s*NDaL0fGW~`=#@+0x%;MEQNRX`4KMDX_WTR$BSy2JfZBfzcxnYAUyCWL zk>QiT?3#9=-(dO%E5@e2FUaE&?5bQjt-*yUGdtL3zbWUFB=TgRw4y|o^PQugG<$k^ zLPqMSoxy>wdxGqtj~``xWzcuEB|Qb8kEBf!!-y4@b3iY0i`pd8Ra!v!oq&7KNQpi` zwv`b_cEDi(;mv>|Bpx>eiJ%J#Wa{_<8L}PgI|ydP*}VtDa7HTDlSX^sN47cy=Vpj+ z18CGRYd2JOyMTPazFXC$4g4QmAp1AWHY==Mk)c zsXbHy=TT5I0q~j2MVF~oN~g|ol(sS)p+}VVBSGSzh9g+#Lh{#x?Pb85w3A2E9s|iB zsaz;q(1gAc- zqPdjVscjVO8>N$!lPV;oRM2{hj!ZvPI-%pn?PyZkWlSZwm_%-L21vRz$?mDWOZ>!M z=V*4L@LK7o$)jz>zyc}&xWn^Db*Bzd-{2RrOpY-_s0K(W?A>hX41K;WP4&aI@VtBo zNe&_MkzGZCT{=aRJ^DYwId(Sc&>8{kwti{D6?LkUDg646B*sv{8bVL%>j zm#-mNAOyFYVgSkgi0b4F@$1s5wV$%nOHmDPu^T0fxnL^fgfr^4N6!fnI&{(Ijw4u8 zY3#1BSWxoQd%YDj@CipUWS9g&G>B&@L>MyHc-l{*gzV2iP=QjSODLwu4%A4HxDV5S z+faubqD8>;y=>T;pH>l50hcCTUVRp$Y3}jMNXVo-DZi@63Z7MT7oeDeDF;db0yKqM zN=JY~$hLH{7g6hJb7&{q3=kQwhBU&TawZEJs?p{5IOPXE0i=FsExb|jvV)>4oyJQk zwAqpOfJy5x#}e5_Y^SU*UiWFD-w8O*l#N`#?P#YQnmT7ZT|>1ZPhw05J;Q1ObIPb) z#lQmmC;U9+dMUx;JtEjj(nj2yqFqy{D5MtP#G0(|ljtiBL?b1>d%Oe<6k28-V=-wq zjXvy*Gck7tvtk+PSQ~JHSRv=9R#ICX0gXflUyhrt+}1P3Frh5O&7Hoc)X1Z31&#u{ zSjQytIS9t54!M0-Ij0?G6)ZBQu5uG&W{SuGJY)Pk2lN8vPtbOMGZSGIryvylt|KOE zR2}3$^Rp*LN#JdqV1EwmBy;G$G36FK$*-r)8m4Kp=P9vBk>rTIE{c!W!0w~k@TK7~ zibUk_rnHx4e1d&J5=0E5LCp2(6T`8jNo9;aH~24fm?m$g&83qph+^X_gb80??N z>2z=t?Y<2%)&p;s<*R`@N8&xYZ7P7;4PpysPw8vrQYD^WBkSkf@$x0ODc>*M(T3V}l6QG<|X&$$>T>IOQ`#y|bb} zGFPF~*_d|9Vw<$ea4WH@SjJe0Y<9sXRN_q#b9acCR%v^9>^U>irkiCVp_x-WI>z{& zpFQ^{n{k_cx}tne{_4}z(7E&(!@aW}*9-!Tn8FesM1Z-!R zc>I$?4?;x?DFLvpB69liS-DVf6spRDh!VvO(76#aeB#8K$x@M8u*w-6I9`I$drkzh>vd_{rxwscC-Ob zUG<(;6*Rg6gbr(^u$&5q(H51o#A@0CWxy)NlRnW1EXB+vBEr|52jI+VcyeGHVk#yW zlbJq=2Z4>JB#e=qt-3Ac%>B`6+vng>Na2?$o=+Mm?h`=h6{p&@9-1LLpHP$Saf)65B9hx@eyMSAW46)2KhJ2d4+mgZddK6h7h4;LM`%2Q4v0ty2q? zR3vK3iP}c2oJc;)SV(L1mXOqkqX}vF<5P@pcLA@#WStC{!i5m$ zsg{T4b0@Vf;?{@|nLikeECs4tSr<6WOiLY(DEHFuhFq{j0Uq1vG>a7`W4z-l&O=1L zJ4z#mWC3y;yv5ESUlU@urZf+4eewuM`Cln_u{&@U%Q*0Q4+x!Dxnb4AaOx5d*@{!D zCQ>nXX7*@W+6WE}cZxZqkxxHat7WhYDCW%CtblCvg}UA02YFaac=HHShGYd z*P?7iUKEB3R0f^cEf18A@LUFpqC3}eS-1Q&@Xn9Tejd^hQti`x2cu+&O|Fcn=BVD_ z)p0AyEp7sOqRJ9%P={dGgO58$mZ*BWu^0)-G&uMFWADv_n!5kLQ4|FcA;`W6(V|d` z5S2|KCMrc##JE%tP*Q7^P0ZIS65)^t2mujNwL%3YDzb?{St1oeA|xTTl!9!Lpk%8e zn>nXKI3-J;kI!@G&b_}ockax+Gxx9GUykD_G3T7mdB0!lmu=}w$BwlPJZ-2?cJ>du ze<|lz&>^8|XXNgyj(*X10XsthGHY2mwi#2v2~m(3MgFXog~V3h}4mrU{Xsi#c4;LuGWQzExdUJ>vA2H};_^vUd8Y zUaa=hH3QD+BvCPZPugGwlnqxMS)IH*$Er*(--4ri+na`@f+!*Y zzJjU}x!9(WO-UcmSLN*mZ&IfsG$w&&0!vxbN%D7;Ox%u zF(qY^O|>T#Wb;WsvYFnLx*POam-+hhT8c9_eK$P!Hc@_T&TrguAWCm`^x&}Pyc5&v zL-9@0Bj#P23GEfoN?1TX=<-Tb1Uz|K)4?ok@(OlxUP`&vi89hH{8M0k>|0x?p zSx{RehVjXji80HDTwm#zKdY1Ljs}6o^a5=3^AX}N6r;NdBEjXWE}y{VIi%x4hV_uS zu9}P6A?`H0s0eS?eES-5N4;wydsxh8S^@#-LP1q4tP-|^N+|-EN>Z}sm=@EYxqR&P zTQ94jkl_}yAR*WNu9H95J`j91DKaVz*1u-l%yM+ zJ<=tW=$Mt=vxNu0q)E#NO@;t)MRmh|BHD1u?J8BJl+A5x`48?6gB6| zpH;z4nB~xcR+!13rSUt#2#I&aWcb#3^CjxLHJ$alHq2j7XsQyD$4cdW*SSX1WJ9g5 zyRfIjRBs<5W6ba#2o=t*{6)V!qM#ug%YM&5eqRs;z1~d*|DmB&Y zOD>>cEfc~XlFs_r6$@;JM)eY2^HV97&AK5jCOWq85wY5n*d*@WFvxX4H?rv&;~Ssp zO7sqpMKhWKlfi76J#9D!YEtg zjz;hP1?79?E7uQmv1q7(yrzYhT~x9li~~den%6#MSXj+ghNV;3waJY|;zQ|v-6K7J z(9f|7Xtufs1VSR~Gxg|Q+%+v@WtA-j0{7hMDxhbx*##HNR9E#jFO@_WGxrE_ix7#p z;e~hd&z+B(Vc$REbcMa{T(apur*8dCA7C+9@F5mx3=_|kKGbgL2cW2HFv>`llSVHF zBNg?faI;*fjpF6gjP-+hV^pe?I!&>heOcEnftc1jPr);%()4_}i2o(tUEW|qHH|J= z?27YxIvy6PU#@bQv~JUK$1nS?(nTd#2&dDj`#o*xpAG%K>uO($>5HVCx0R9#sbd-unD-w9k zAUlPZK{p4(OB!El!ZB2an%)ItceGO?*>swSVRwWmc^T)Ef90gx-l>e5UQEfIqpT!m z#w6-U05WN+zG10$EeFPSIW>%GgsBS|w6Ee9S4q?v;A@rU0QC^Ad&JR@nzc{=a=;Z`9;Ew3PF z`F*-vnm@Pen!a5O$v0i zoVOaH^7LKsZnTSTKV|!Qq_jz{l=k_K)FdmSu?{}OFB7Fpc;|S9SbP1J6Cg8dL8Nkq zlJbQSpc-}T;Rn_-e7t);v|T~Zl@zG*z8~U%?hZ5ReiOVk6qc5DrF!ymgriq}Rqv|o zOZi_n7j7Gz4c0Xf&wx$5No$8-x>%0-fyLH*y}4dMZV8#WQ~73^wOsG4^^3;`AAp12Ux&h5)n&yw#&2++sWE88^8H@V0UxSX+ z273U(f_AiJ{v-Z$b02Kp_%_F|8HrmnyvUD4K%`!-u{@PF!fVN4uJ^;jtT(S;mh1Z* zt7aH#6Ro@KwU%|YxUIgXmJDql+B5-=jNMIQ&mQ~~&PKe-t@Ybp%~$*ALJuM~MGqfM zTD?xUg87QR*(6N@^dGpo-

G6cNd@Tfr!o*hO=PSu+~AbRWt@Ocg1e6r)7Cou|9_ zBTz+c1)9iiOfJo>9X?r>|6q(Nhe$>Z4R6PKYNEzZPNlNDA1a9!oQ)d78(Oa|)_xYc zQWrxT#9D%biIuBX0|tB(yTN=M)Y0;q#|z{gUW;Rb7@z*x76xL+;9RNK?i82x@%xB* zXZzQ|3R^)miq=HMtT$+`5o;0dbUx$se0W|LOyoG2J zH|~Y^4&@kbIis^_rs0qmeIde0t=8B2JvjcO){F3k|c=9QPF2{ggYid^4Gp-F% zcC&WsXxJm_SZ9rVtaRUzso#aR0|U`ld~tcKXEe5E3n{%;`&Bwg9cPs|ZvRS;8>hXb zn{NP-iFNp+%D1NlKwSsG38t(3tG7@S-CzWV&oK%2r}t$I5Kk@R$O3 zmG(udKQ# z9phNm$A9>b{F<7u5ks@vIQh zs0+mBiGb)?j=CN~V!P{NJ9%b=tIH9sblswj4WzCewbI70@HZF+y(LAk!&OPTO0#<% z4G!`bdN1@j@YkU3O8?G=Y_b)>gqkI@W8#kz1z2;mFAAkAW>dPy9pRK&r_+j3T|^kt zT`T9UO0?x}A3Q(9p60EIN+>zEc(#tB1{X53KpOxL#*eQ<=jF~Ra?b|7RFbB?glZ}i zq>Z>49h@)M-$4u@w^szucrn(9yZ%Fo@pXdP=b$LzGrjloAy4+?Q9;0c{pW6T$Az#m zNtoZn0RvR%iy5mGNDXLy&{uHEHJUqmTYa3N#1#-WYfoZLX)B%*jpwAfWyS+@jS7_o z8O`q`7&Z@RZRAW-6r#o+8^Sv$)_ed_PJiCTW+%sMeJksV@-9*a^+Tz}vl)m#AAt6= zdLJ}uTHxZ1u=u@+V1da6wt1- z6Hf%!KGN!BJ4_PlqUH5Wc?Wra>j3-sINe%Gtqs~}m-EixpM3TAS!xw}OhK}mk~WFJ zw^FF4ZA5Rwt|tfQr&&uNXEEp0lWDpEKIKX9sT!0?3V~@nUq-g?j+8E>vqBr(f22Dp z@MQnJ4}|qiLo;XM=6zK5w6Rq349kV(-bJ>UCN}*EiZZH#z)nuFqxHmCBZUl^nQONm zH|)c7h<90{dkBU)Uc?9&J4(GRYEF=xt)R#$F7Zy#x5*x%-tl#P)gq#lwglX6<!+EUCAA`zVUSR$#Kl)5xT5gZy*NfcJYf6g*F0N<~T7^qlwk7 zyXgkEX3SwQWNnP@F+FiMMSTqxGuOjK`AXu9nPCW!CChWQp}PFD3xqYu9;W$ic{)hi znCl3P{L%1J1D|M5@S&E2f)(w>ir15LE3SvIoN6TpSsFW@YTaRf!CRuBjhL-p6UQ@z z%w#nLn`sw}7}}`ZD1H&JCr=(|kApm)Bu3qz?|AlX@D2LByGi<(W<3aA$u+bu80uTd zAD`WT`<_q-`CtTTl3oQYS`WcjcLVROsAgjfSZ2_%XXv^oXtX+3tT#<4Glm(NdQ3cd z)qHM?XVBNMS5{KR{q&EUKAtzPX>gKe5WizFql^LhM$&$TN)>f>kz&qv84y;k4yubp z6GH7dN8_ug4V_$9zysvs-QB7@x%ar9@FQw2THaN;_&)-Rb0=>s0r|Jy)IU6Q04hDm!K(~W$QX%|Ct^o55RU-w2jnboGn$_ zctkoXGZyL{MzsE>XvxI>5#L#0QI+O;;Ws2M|4IA+i{~RD>es*ON@-tAwGsIqv)+KP zJh!O2rgsA^+Kibx8O|)YI$rON-l;^VqfiioF7XHG!pgbc%ROCUspaDVFQVLl;w*_n z_DVvo)y@5?-C9D=sjPKEO< zs)pt;STIAk{grB?V)&?*C61&r^B`|o=`2w=k{W&!?=Hgc(4+wD3b>+oAg-bShOnpa zD%TK%U7BfpK?XywPkt|*AKtGQ`16_pU8OPx$DnXCHRf-Ss}97>*~-3EBG8eUlKQN; zZoed>pKTVS32QN!f|h7NBYhN}tPI;jVYsY9`U;^?1u2JYn4HV;kkel`$eH<<#z@(` z6~0)jDONM<2HsWXjFX*x_z*qKrx}iLjy;_Mw__}DW`?p)?%ClLdNO=2F2Lt9)c#!&-CF8*DuC$%e-%~` zaQ6OF>I6U4?t05xtRVHytrH1FM_U#iP0zc(V701d>RGF`fpQGWC2ZBPX6R?Z*5sj1 z2VN(h6|EPG0kJh0W~;P$cp~5I&(;biMh{pM1RSf*{0nts3A`21X*tG;`aAdu&zhh1 z)PG35Co#Gvsd{hp5QKXCKpH(NBHDb^k`+l#>VkG%op?qV5t%651F`8dW&ON70Om0^j8U}mQI9eIz+ z5gjA*3Z219b+F;ehVIhJH~clU6}pPgQj&ZN?QVpZ?fXqT(wD~WKH|By;u}AeqT;^(Cz_xOG-pX3lehO=eYJ_!`O?)u%uoZ zpXW?)0klO6iNOK1|4xJi_PsKW%?yl zDW3^v$s_n!1GGGdgDR2tqyou=Xh&$5x{cQ@U2+_FcjgE<^fxir{GM0Ulc^Hp?70FX z4WDMS_rb5v4E||;5kS&fr|j3=xFlurFf$yo|ID!5qaaQ0b>FKPE|Q)oCBes-*LVrF zWE<*Bb|k?m!L}LN+&>EpCM-!xr$i-il3MRC!R1323CxQZlb(P{_ZE8eiLMpD@AFcv zE7e|1mq`m(Po|i<-Be&0i(r;Y)Ob@3yToK?h^QcCw6%7std6q^+lsrwKsvnbU?@sL zJjk%~yYV_jddi_!ULWnk3#Twvaf>dkr#|UgyGd7r)WK-821t3bluLWgX!eXkG}J>G z2iPXBuaUMUJE{=xF*vOB^Hx557yj}hC(~ngPePb-O z4p2=;!3;>rfwos^4;KP|CmSTu6w-_d{!OqpoPdSxKnwV!e${k<^Q%GVR0FrTVSM`X zqP4p1@*PUrcT0SWv$*S^uVL)97MwlG9heAe zf+@npQi^SY)a-nVF6hv!OWqf2_HACwQV_K^rm~6rOaB@;#zhgbks?4?yP=;a{b2Zp zqx1Va@@+o+2AV_)>}|O~cc)U#ctYtH6O!;-PWv;A3qKz7@|xQbyBFk_kgfUgN$Fo2 zPmpIxnr}iiyFSpHReT;mVm#sEF4AbZ98WjZuc4mqE+gbCNVzeIvp)ca9XyD>bgNx* z(sGH@?B)2mblbkYgLvYp$8j#UaB=qB4uV@4Av6!JeYd!nd6lpYWo+8_x$fbw0*rRQ zxx?^5q8vNW{hz(t`cnD+e)>zXvlgFQvwE5Tv+iN$`>X#6lr1`rqt|k!R?xvno4kOX z>#Nsi$aR}De*1v&&3P3(5u&O7*I@VnO)-x##~d8pHurMg$6)k)((hNj*17+24cQc% z;qyJKqu*fPj;rgDBk51Z^q*S*cz~xyz8QU#|E=x~efkvkLaIagvP&O!tz9hqitL8_ z=N(oD^W;0Neg2V!?mhNC&;H8Y@S8R{XMWgfx^(Lb*FuKxcY&q<8ML~v!f5B7Gt8TR z3APf>@7d+x-02(~R2AYB*%jsU9rVYoC;N2Mh5H;Pg~toOY&`%p_8ScTm-=tbg_rJA z`=;WSRKAz~c|-ST&5C!+Kxdd`|Bohb|F@>JuUq}z?XWwyVqowp(0w$uL2)T3m$`oVSARqI6# zZeP9(w7d0T_0q{h^u?BtklhZ>4tvj>`PLzF?^i%^ z=}`rw*?;;TBio?)312-j8o^1wPrTJw|Ih0u8_#}~pMB(yoy+e1J@hV|()LvPF_ZqI z&ywN)+Eeg<$&erZNwCg2^%B=qHpZwq7_J7}DpZj3`KRi{D0Y$v}#C3MzBgtW5 zu1UxSoo@Ke1^++%gV+51+s%7{%$pH8r?es4D;+}IHX6)-$^Yki;nkI2KKlRh1*2j| z)B6P2g+-$NtLyat{pfsj)#@cFlCIK%NWve=b2e-i?7Qu@teId7{+TWOd7_?bFei2R*WfpH4SkHIFP4bsHp-;0oz$2%(rEi$ zq5;8=I(=YRYXdiTw@TN{YOzQ$pSYA8u&-YK$2bnXsXS?Q0&N2v@B9AQdI~<)rr^(P z(Aoh+p1{Jfl#KdE!MOZV!FN*&Cz|!C|3!122B4H|{fa4;1yNR_U(*4?=X??D9-yNR zmfu0MWn|-qZfDmMeyOTiGbB{**Qy{gipRI0b}|s@KkhPc3B-XCuR%LOUZybuAg$G# z*g6uw#P2ygMKY&mTqPbP1kyW+^zQ-NCuzOTHPQcO3ED%;Zo|xOe^6O;`x4MxwRkmS zg=C9KB0%52pCukcesb)-jTq5$-%rRFGQr&P`^QiG7%XB|cSp}URX$x@WZ5Z~t{x|A zdnoX9MS<*k(~CFRq$k4bjoa6uA??;j`xz^ff$wx>Yybu4TtTbsEJ0 zis=71f6_xueg11p6H^ERiDN6Lnx~3EplA*LjQr(XV3Ly(V@&V-ibg~~DQsBT1Sjoo zG#I3hEYOYSNd;<_Bi9+!xg!mH59TWHVz*(lGNa1!Dwrl(+hN)UfNr;ogaQuL-ADpi0qOph+TO$#6(uhdT21r(Joje2lBE%yb1tZM?vP@m+YS$`z4 zQ@RqeLxSXe{I}lGgY|fC8SciSa~}vnL<0sR;r9jgAd*oa*`Tj_I#p!$IKlDG+oD?~ zf7YZmggy+aJdxN;KNp%>ySN;lkk3!d25L5g8{7L5IoDsmO$_)>Lv=sOhf|ay<}3Ew z&>+-RhRxHj(=l6^D<@snAqU{`NwT?5L0VYH&@>2dc{y8B19 zd(m{b87OLL&S6UF8f_T(Y$)q%7zEpFu*=nY0_5Z%yNw4o{d{z$O7SdH-mnx(0(b0b8D*vuu>*klTe3jff-Qi~ ztdiu^AQ%5*J`cz z0BwtQ(-7B(nyx(t!cst{=NVu__;7x{$Fwe5KB1nyj{ZK+xo3gP!w6Ylzz zR*$_2sou~`Q*VYtHF>rA&seGn6pcj8Okg{|y=7kx4JslS4b=ya`uH;^xt180?k$SJ!j(4>vdsPZsUbaVy}WjjeoSg3EAVdu95R9S(CKM zaFS(B1JBV?pp|YIi@{x`H+xRZHYivio|?j^XGa;o0BA1Jky6iHu60L<861+L&mJv| zR$U`(?gyA_Y(FRcObQ$$?k0Hh#2kNlqSVN*UD1;ib28qQQwOqVZPCNBK6Hb=Q4@;e zgJ8-VWvkFg%?9Wc%n+)UL~H-k)bBf~-%3^P@As`UjzA^~xnCv#K_;G}bjBSN*ZC5y zP`V6kRrF<;tnn`gTYBw|A%z`&k46S*j%fSt+b_6V5ps4CKJ+n_sbAeeiGg1tq-&KU zxB26)t&MGM1F=JMA>vCc^Xb`#(}&jIjuO>StH?SK!Hsty`CHeJk;10RB4z~1MQJw$RE~53~m?XOLyjNx6 zm*c?(IzQdh&-AIpTV1GoL_AXo(=)Y^h_Qgz0xV0ln)T5D7dy+Ykw_+4w%S-3dDLHS zdk2Z~0PoMu!#*mrsj~HO)peiVaIuIkp}B^I>>VCoSO2qDa-eIF5!h$H}hs=OUA!_wD-&{zX-};B5@Q0 zGRc{G0u*vvPVZLbNXXsNQI?Dau6e$G^9baggq#i<<~YlFn5`gDt^Up(*83-DodFq{G_5xOg9Z0u+6P4Z|OmVS>Zw@8> znW5`Ym)>_#egl%$)FY^~jAEe&9w<}V7PLy1E@g0nV2$2Hr7@6dak`3Vp^MM6W<*th z4^b|M{jYPEyE7=hRp`aqo^tJW&Jbbk&zDG;UCKIRjm1>k#|D1}wCizH$V7}MP?-CN z0RpdCi`V^3JWG1r_8w~nUaya#9S_i2U@Q=jBw0d2zzD}{5nxkbR0M7+Hfa@f^w|i0 zxp9}-QGiQ>T%SnUv}z^)q2HbIH`?9Lk>G3o>Q?ef=%}uUU<2ZBfqml)Q9{5Jj@u!O zT*8(&(5@j)7@YkDT37XU>U0;u3oc?*qbV=i?ua@QvF4MG6W7iq2Nw7%E>|fU>~}qm zNqF|3H_)ab=t@P4mp}+W8oC>dP?n7NqK+uHWrN89BlO%o zwjT+}VJAglZ|HHqmUV+eSU^Yj&M~dV(9|8cpnK?TVNuDw@qHQ+HB`pPCxR*0kot*? zdOVosS@X`4QTGL{RiGmLzH&!IH_@tYmbNx7W*hPs;Icwt5Q*ymPZV%=Jl-Wsduw{E zscZ+iYc(372rGM3(o)0?$Z_iM+kyS$E-c~AdV@Jt4ti85YAO!cLCC9H?wllsBOS$j zEQ98W)ZeavcB1rz$r|kzv_i==@E773&w#WeMWMOi6J;}zHQyeX9{x3V5Rf7KA3W8` zEYzJR=vzkX19j#98su>V&o|6Y}uU6W#*n#pvX<}@ z!Tuud29`wqZ#17~FaiiLYtUzZkM|rwH84+gQn`IV6J zX`2RHT7CsWE>+IfL>C`1-Vinvwh}hLcZ5~B6s@OxjJGm=G3Ca2?FE;+byQ|2cMZIS zJ9NnKqX+ACimaGVN8Nn@XGs@GuU>9jDAXFm3qTI?@}h!h9*TV0-4BtIloZKDz-yXw z;$fypYU#_$H;EUmd=jz$KTszRQ8h@sG@nErLDXpu_ z8rZ%J;nJg-6N*L<*3zziq-!DAe<*W!jxrVLu?lRbbK?G1^6NfXOu+*X<4U}mOthm` z0!I%4kMa~dT=}3tyGlTIIPvjGQ(0Z9WPU>q!5Vav5ge)5a5SYNtO7ONS=zaElQf|}IS?I>zH30Z zgq*uq+=zjFl+XKhMgL;}bWI2R4Q}h)YY681|L0O+8 zUSI0hhsr&2o|EmUq5BbBVZHWPe8s+paV3SVn;NIuMrtPQ=#A1D&sy5=SVR0JpiRPk zU?{{4f(3Oo;$0DUA=!k6nPyq=;?qcmKvFLkGYK0LdG;)MPC+G@jVddW;qftk z1kdC&daF%c44e|VQH*>}aMiSb)Yl;TiPub?X@@72eVrIdxkAb$#*vMoO^9758AEkF zF|dgfeF~u*W`E*DBd5X2U(Y$Jtb*-2>3OaGdb{CM$mMbn%`n#0DvsI86wyZ|jNG}f zz!!)#ao}>S-sHAszl7EkWjuTIsT3vir7OV&!+@vEul!U^BotkIP;U_rg5;)8-lRvy89Qg1Ny2)T$_$9)R@;%fS5y6YTDdIZ^_Hkyr8 zCqBJdCXOUZPmfEf z-+9!GC_b|IEOdi7cN2K2HH=DScC*`&Vma$-9X9~I9o`Dh@~+`QGrB!+j1bLzJX1w{ z7h6P>Q4a;e6Wv_rLBP0Y`VOmTCe!+Wbkn_JAFN4YYSD8$MiPcGc=8|V3-T{20WPnJ zN(H{0=P*MQvl`9NglRYHZu#IGK8U5Zc%r-x6ekb;%qACH6)|l1sM_Z_ zdeZI&C+$LVk|web`E`nDL)$2giKONq9P;V)|3+8T0=!ZsmowCP5Vn#kQe4_BncU}6 z`uGLm5lqgcpTT6YT3ct*8niyZ$zlf1*?pz@l_r`9+HZD1pG{i_2K}7dn^=a>71$M+ zfsyZi52Guh-$X*e?WDYgD3E_!^L83*=5_Y&`3X4mgy&N=`5wqrHlW9m^T~}1Got!w zsRJ*Tj)As`=cRdV#2KE37!}`2utXaXl43(u?i_I`_2b@_K5%)BBlZG4zS;RtyWsOT z?QZBdQMIJ>@AZ!Tvz&cq?_Lgg)F2}5>SlE#j1iI;#Fg{+L*+9IL`y4Wz-zRro3A_$ zf4p4{j!58!5bSWCyB{j+zK_^-mt5_?d;;~+Zh*0ZG)(Tyc_4b>BI~2hHu~LLGN+oU zkx7@nrJhz8cj6>uwf7*$<_vOCzY9D0HTK9-xkIh_=TG5<<<;?^A6q+Su zt!9(*ivuRKU-&O_J04CF@JQKnc~3`(P`&ta5Oj*fmVda3Qm*s9XXT{QwN`1s*s2YL zbEZU!2!7aUauc%9P0G{T02P>jlCtzTFJ6U)ySR+dV~N;!!nQWuUFb+rsw;5Nbn{k1 zCxIO<&#$-$MAeq%)xQOL#e+E$e>W>qG8lZ43y^NH;5$@Gcdxwm-Qat~|5}|kI`+Yo z9rjFX+55m~tP3@w)5UwR{?c^bO_8FYmme*P6I(eQ zue_xQJX;h*S@UWA;ogAV$GdO!tl9jy=_|&~b$Q2cIvOmQBHI#-p(Ch^)Z9wJ%mxRa z9muCN0kUk7{O^AalBg+Bgth2N#N1<4>#WEEB(J)mmJX(|bHPC)+#vt;RiFH)b4uT>+q5e_6V_?iMo?!vyecAKOehV zdmMcSGXzDv)p#4kY40oAX0&ol8Nw%nGaV++0OS=Urq!NFi&fHr(^vk!IZV#sM?IQ6;5@h_KQ{j9u zWI(%{5vA|BzIo+>^p`5Wk?ua8f-H0q48__rfb<3p2^0a&I-cb``w@zbm7681AHY1V zwV?oO;?w?VeaWNv%MzXHg(sn#j_!2C&K!GnzK(O$EoZTADC&)#rNl65G(p-6;H3T* z2-kaHzLGJ!8V(WxkTjcPqG+(9T6NJa_3klfgQ5yR5)w<*@qRC0jjN~9g?caG526+a zHzO2*Y8e{cok|guGMB`q&a`a_Ev*ZP=dHw(tAvlJ+)ePoI|!4%g|EOE20tK^&#}>$aIM^yN3WDdQ8PtgXt)7+R3?s_A=b7Fe}CYv zZ%Z?|vy-?}YpVGIWd-Se0r#Xdi?CdGZ$0A7%mHXR<$Cj^W8aGv;Ffec4szv-McBrM z>sm4ZI{hTHR^3naL)fLs+ogHV}-A0^=r{q!Ltz)g%B}n zy`b(h3#$CE*m+|r_1wmCAdDJ$TdcKgmK!tl+w^~8Z2+3ouD;L?R%eH&dN4LNo@5lX z^nh?FO_vm|;aGu4Wu+<`A|t=Q|NH>A%V!UUl9vEn+3_cwj-BRz&8JoDirD4NoS*Bp z{Mwl*V!b}CzKVjRpHZdk5Bh3Q4zdRNc90m?H)i^x?R}{xfa<~rh#d{n2^>X~&ZIQO zPcgnAp_DW_{X#Otv%O7*^$(L21tCFjY(YII`n=Q(BFbvm)Oyeg3-;Q0-e|dkHNybG z_i~)jzD_3Y1G*zc2QGZTQ+v2_GG;woTyW_)dk5~=;0))CTxyJVb#45vY_CcVwkY~WOGR`OC>Uu!31U%1FOrZxy~Bhe z1<1rfH%WTsbGTJefAbMF?RXm+4fbm5sqrO)o^wR&xXUqaIr?$RGL{21HB@Da-rjxk zyyVkC*0)BWrW%AE1pT6`qHTk=qrUK1=azW72>}>lT6ovB-ZIZ<5y9qjuYAgW;O)ZL zfYe6G#5G<#%W;bA0P=*!(oUxKMIj+w7`sHEJIs*Jo<3gs^{W~7Nk1tG;5L5K^I0#j zE4{SNns}O&-c0H0%RNLTcST)h0)Nj}tS(2efN4qcJ}2k|5W-y1rYorswo1DPjB8!k zydrgqie5O{ zU4`PSsaXeMT()W*dN*B(v3o;}?hcPxA^x^$*aw#*;dRv+9$^pFQ+z!DVJTHmvhm02 zs3`!N>NVayys7R~;*XethCF&B_za$=YLV_1%_yGOv_W62*-RZ&=je(FE|XHWH?>uH zM#!2yC6ng1k+a>Abz3ql){ok@||v{lEcfk$A=U;?B1w>HMDbDUQ7wtmhA zO(<=Mh6kR@6y0@bmrTV7NST?{1XHOoz~Wjw8>Zk{Lj-XZ>Vc?^05qc`qFZ7o5EY8L zS+?J2@ukKf&z_p;0sej&F*A(38}7S$j*Yb>t(>9t^WeQ;lwz@6wJu=lO@bFhEDS_K z78JZr)rtoQKZNt?ILg29 zv}|)zfaxcmIZyT@&!$Bx!FAEi>2VlmXcKx7e#c4DY<~R~Efz$H1Pw+fpojrpJl2AG zeXI*$`(cl#0ih&AdW!|UDI_!?K&1{cp@@+E6*SmVi4{vb%9(pyQ!|EU&Bn5e`dgvv0&6enTEBQL{9V-zj(~6)L2Jp7W3F+|_zTl6faRCqzAs7Zp-fwVrxHMJMNfaDA{*wiZNS|M zdDOhcDXY5s%pxHo+JR8y-0X|hpj#Y+92MyuUnQwN_Wse`15)3IeVj~7yYiNEB>9tS z+sWh<=N7lK%dQppb@LinEwMA-flVTh*ffztbJ2ba$BHmxR_Z_Cky_9N5Q^q3GK3U! z>i&sc2ubUSwVN7ug%|h?-CeD71vNrmR`3}8>hHh(*dz1w`~Vej@;yuPN0zJz3q3v<|;gh1n0kQ#ltZMn}i6OVHNNlnBR;KWkI~=K*PSm<+5n{ZyZY zzTf2wbfZ^|0%H%b-xfd)PpI!aLko~Sp_mMolinWHcS~i(FX+D{aD9iNjj6!Z)HGRR zQ!MBq*+K`AHYoTcB^dd-PQD}j12mZiK6(K-K5_OSz+>Wu_9)!*qRkkT=V46=9@9Jv zD6|SnPBK~mfVO?Ib0ToAG$+btjb%B{k=o$aXTy-UA`RmZretOlRFCZn z|E&4#mGCYLunAWcyf_}xPRla?b z=yx%WR^P`U&l|MN;nWQHHcAqPmn4@lC@G2ZV1nK7U+rbl&Q3w;gpgrp?92`Y zCO5`Kmp4bwZWCuD^Bie|JafGtS{fwAiu69Edl6tJ1Sco+*ll^S)|IdUCCW6crBoHE z4P^P*ACHVL5eqOzV8&lH(pCnP>fZd4pHlA-KWE*7a3raq0aMC%mf)gN3|kyej_04!f*D|uyaN>!XK&3D+BXW_+?G0my&Laya7 z;N#zJojA-HY>{D_wRRdu{qJq9Ghqr`9?cmx=OIDCw`Gjkn~0-E{c?m2d>aHV3ff|1 zq5_{ftPI=b35-&=5kmpl6xwkGA%`$S?||~$79{qD7HGRo*8*)avdXUv!?ie1dxLF? zo!KI*3lzEThKt%YIFQvMnoW`ezHcqclvSZw$R{!qo1>*@Hc6pV&xVp&1^|(HUOxFl zf||AZlEYlI6ELggci9%vtdXTXHLnv(igMkC=-WY1RbbcnY8?v@ly#g2Z8H8d zgPP;eS3vXac^9Gw{B#h7YzjNIUem8fC&H0pK_JUAiW&$6@+b*b3hAH-vUd6i z8pLJG-qh28ZJn=k8WGq|ZG?Q`3PwNc%3yi;U`8k1g0PVq8-$ecV0ttNPM+?@nqQ)}oc9TB4;a>8H!DI4=&@lJ)s+LB zSh#*nZU;J!oq9hL1&CW~(i7CbJH~n#3ap^g<>Bq+B6(VV4u~WC7H!@&M6(@@H+zCo zP7K5j*FG@+AEJb8-oit$S~{=Z_E-dHWZ~j9KbE#%tzlx8${kF_?sJ|NQRs?Z1S->EIXL zMke^oi&<9wid2g9G^c>%8DMH*0 zGlVJZPUwHI_wM0Pzw6uZh-@Nbmpwz2N|If6V}>M2LMlZ}g^(m66T{e)JwqymjL4=W zZEUigFose#5k(nhlwFu1d@wV;ck5a2`u1DL^8Kx~p67ku_j!-?hvTSZeD3>m-J9#Y z&g(pNJ|h(hXOyb$6^}l;#Rv574PNae`*XTqkEPt2kn|VWDkF1LuFXV1e%sP$G{Mws z-Mi9MC9B;X`;Vjc>vWMwVfD3RxjnFLO%Z&#FnHUMZqxOf#$^djWpOoNzZPPh=|w67 zO)^a;Xz(&0rV3-}Z71u{#n=~3l{LPLwIe;D+;1wK-rY(mE28cq7HgEg`+9Tws++>Q ziu$)tm)S6hBjORZ1y{m(AWxqxZU zVyO?==~d70-N1KJHrH}%ddo5e4ZdX#eSEi=n1wBX(Y@Z8R|U%$QbJRVFT=y{jB%&2 zw$?jM4Led25#m_xMo0Vxl(IP;ozQ^j!Z|W}ntVq#2C*#Z)Q#AOGjv8r^@xV(PKYT7-OE}(B zPLduEu0duNW~kJu9C=^#GRf(DpMbwa-;FQc{ZhNlLf5z4e81&G9fLyR2AWl*#Yiae z49HTdoP_T_bnyc1_InD-ux0ic`x(R-q}H(QhABl|rehq-=5ez(ZorlBU|d%6zQ0Cy z#hp6Z(0hp+-@ZwIpz+G8^Sz#@_N=JNOVzw@Z^nGU;16wm0ReE5D`UpvvcGGW4X6Kw7j0W;4r%Qpm3 zJvLt!0(q~hae5FkIH@PNeUJ@9iUu?fc%aK)9|vK?ARPeGF(vPAWS!-$e)=N)(DQ2F z4({yLyh+d!L!xxux3q#xCFZI1Mtdo$%I#v)POg!u(Iz zR8bd~g@oNd{@OH+7bba@B~T@B+)PUv={ieajkRMosvFw;>x>XtIO%wZLAZvYsS%Kd zjM=8!87uX#<_dZ%6hg_1!Q2O*JC`0b+0g9V@6@wuY+0YV8!nO8kjE}YyK5#=pI@vU zamK`HL7rqF9F&&VPV>8lmgPuKc~bD}8?>PF9|5iI1%uvdfAXVw3BbN=@NT^IdFLgK zY!i@1D=`^cYwvsatxyP{I@Decv$qKIUERSU6A}0ejZ7>{s9RUz6iy5Zz6wl@-<4sr z4lxFH#h>rte;imHv|UX7ou{3%^Tf`PyVE-*jpd@l9-fq}bdG@EI4|H!yCd?3`^qKQ z>97Fm_~jiT+qm!LNJcK|Ds<GAVozN2BdBFPx(025qex%{AxFgIo6)R_vV6I zpdruIPpa9cN8R|MQ%!kt4u-V?#2?U2I{2Etmd5UVC?pl5%RK zyqi*U^jyk*4?e!)Hoji&BUc&tPA}3d4WDc&)D?G@_eA3_8SiaM*I`itpr1toW~p*9 zYe`$aO#s&1I&)Uq@TSnPXHt@A$7Q}soc!5-MWweF?*IaIi*Ps~P%EqicNj5q{AvF0 z_`TzA?*`J|ZFPjL%4wY{N=YfhDS7oThDYO;6qP=XL$QZC`Zgxjb9AEov*p7!!a{y& z0FFBKlBPR3NUbo8=hh30C& zPuTlaCq@WtwPYk?dKu`WQ6pH&jLcsNh%hdIL2L%7j5)*Y1Vb(E9CRSAbDg!ROL#NeTxma^STG@LJm_p4m@)hCAmbfhcOSlLED_-n7D@)RHlTkn z1W;lwq~59l%RQ0&6z72rN3u*9HnB~~=vP%#QE?lny{I)G&?BuFZm*Uy;fW}TR>r`L z%aQj}_Wi^C7k5dtFpf!Zt@!wXN~ZygRj!o}cg@#6)l zsS1!gj%W1Ka43zeM$R_IcUMhD1|hIo4bOwEYD8`AsYn*20BgIx>iqKqwsDUyma09! zz0RY&t|rQLX)%Ib9=?s!y+l$dWBGz)Ja%9ou^d4`U&qVUc}QIgB#KNhyj^!d=792x z7RKomLdQ*UZJK37S|Qg!Buc9Aix4aYcy) zG<-a>po36+3JPeJ6~6yuK%?Zma!$>G}kH`CbmI^S=2OaihiiBUz!b7+IVN^YF8RhJuoNf+yZ* zPD)lTA%yUaW*lX#4`@NILy6-y^_KnU{L}&eWKG*bULpIiQF0Me{Wt>-;o!@`d$YBZeGBP|CEMw0mgamdZTnyjDz4bby zf6*I=P|*-+u$a6%#A0)nUBA?!DTnPRe zY7d05>RkYMkgR1?;63BL#OB+In|gv)tuCH{q6?tzGw);WIt?y%eNcj~ubM$`FcDO4 zO`Q7bVLHCnI)MbLTabNPjpG_h%B58o2OHCPgd+aFz0E*jh#I zBEmKnF{K9>Qo5)8q$KMz-{)Se73?nBNGh= zmtZ1usdd`$N2m0JqDaqaJn+5KoEE|~?~w(Xwv5d1_f;79E>=)M%m<{|<097Pl|#gh zx?33c(vl|o1cFO8x;gh}hVFT~Oi6$(Q6ov^!X=FTowUM;pPt`78TnR~ z4{6dhuQedf*94{^hj{!%5$gw`fzFZXuS4;L(Czl4i3j@23Gv0b@Gd{l$JKZcg@Tr7 z)K{bw*Ui5w7mg4pu(2#^fCLdjbfRG+soTAWC5yE6qm{vNJiWewWugO4=YgWBAXe|! z1zq7bfM;V^XqrZVc?C27gsu0cTcijo!`^RGct_2*XL#_TKyW=9uzVPmz|ZU|Xe*c^ z+u?%xS>lXD8t5Ar@ZktNWfG#C?Vu^mq=oH4_uwr?qbroVOm2#+UXOCWoaG0WiEnOI zY-X}ygQozXM7 zGP8QDSGT17o;W$vH-DVZru)aW$_m1pUhyN@><6L5M3Ue3IzMJ2AUQ06evsi?WeMf8 zBG4EmM~%=NuElYE%O1`z{BqJe#P8KjJHBz>{^^PlYCK*9eW5djTCUlo$p^~2?W_B6f}mbu%Xr_; z5%t-Kd4&*U+2BwJ=vc$H6Q8kDCz$AD-7Vv+~0NRSJX6KT=;06gQ_I zop(xPiI1=Y>=~5$JZ!q<4e201`1zHK)dlVTK0aZ7lKBsY%<&$oi;mAuogCVmv4^Q! zS#`r#&|Z@G(e*J@@jxQ^zwniaq-Ydpw){ zK+&0E8{40=@nvs=IExp$<-zDfez~7CEU|F)u zkXl8+UxT<2n5XdBR6#H3q!go4Y6QH+`JE@3Dh6bN2T(-1g4@YnAlTyoB{99zTk}qs z=R(2eiI&|gtwN_yUYQoD;Y`>x0jOo_jON3>XxpBhIOXAK3{A2A z!KcBAZs{A}Ts&!nWN;Fl<#v1%$*n925y)ZD z9MS^YktXm(_pw#L0{>!4XB0*hr870)7$@GE>3E%^!|7_!2A(YhK9+HP4=$uZyqd8T zhf~*#`mxH!kU8tJ$OD^3IeXr-^zq&vPd;O80nS)%J3Y;GSWN7`pzCYiND@LxUnNO{ zjypNZ3UENvY#A__C3g~Jza>7vhRoB75~avPx=i0&DRnjHQdFyY>SO>8z0s^>3bm!r zt$Y1Um5v2t0ptc;h2hkJ6u>OPt}bT#eZnt#k^pxX^k01?COA!E2AtsVQ$Bdfa(o=< zIjpCv+sGXR<)Us+P!9kq$Z1a|WT(p=(ToHyduML|w!=Xg``tP9FNXmq^*!1b@CIy$ z5iOmcRzZ9;Qm+%?EJNlK6URf$Cnfi%fLQH99!ItTiZ8_N>j_MR^8k}2OcDwxaQR}L zY@`l=D`>dtWCPEKhV2s^ZjdZYAie!*0u=~xD$5##ak?c)Y!0$>>A?07>5AOKxH`rx zn}!21t0d;}M9tVbkFAkUe26kAu}K6E3qLG4wV%0M>laf)iqv!2-)CZE46#mn6n{=W z*iX^w^z+h~WPw6)R0>Fi1aX+)GT`Rs!_Gl8ThThB0i8&PG1Vf4FIP*#>5sq@(Oe6f za_50e?e+;HYFxO<5X-cN>5|AYV`{{d`IJC!z|oZ1m!_AK5is(g;`w#UF$D>(KI{_3 zh)4!vEj={hs`b~Oyr7p!VhI$Z`pR3I$g&=lMYOQ9^UkE`!!a|HE0c2k7t0D#I@GfU z+vBGDbTcPurte9;vTD=q%m6dM+?JvR6-G1WfUqPU+CX}7F*euLi=&*?1WHSesAT%A z@KlH+*?JT6R7lLjV44db%{D%;{_Je?K-u;81}VhgeCG+-dp9uplt8fc`!$iEZK!;P zBV7mQ2kGR{wRqFw9Ax_ysQ@u!pU)rI%{9^gBJ) z0$ryh(IDz3tQ2{V6Fm0ESH_1JqHpoQWP?;br`L5`1*ou_4_n69@iF>aIl@4-RuP?a zygx_f8S{2xDXr5`x`f9LzTA9WuKl}DE zRma{GapgR|DkA~WOnf03JA#NLt_zo&L_b$cdd%`;U?__44Enn^1wExdL67c@73C_Nso4+YTrDP@WJ7(ouN z&wJK|J`~T?3bA<^Fd`bHA^#D<*-v?7Z{qquP+s@S;R5Y@8Yi&kO-3Qj9_UD%Z5fo+ zE&I-gzN2XT3GVcvoqVJggWv585DgHL-(B{s;jwlRZZrz7dhD0OJ?P@nEkn8Lsy zR*)7}S;|y9!RSvxNxEP|X}qx|-bO{yGvPb2Zxnz${gnc`-^NKvsb?p!wdPEdM>Yzr zaX?_M0CW*5s`er6M@mc-SMB;Za=6owbC7VgUub|Q0#q%*uOl|jZgO>QJyYzRUihpC z@@zG#Sc;udeR=y$ze`mQFn~ZG>lvx;*~J=2fU$Q0I!4re+%ena6}4&ZGL{58hjoSy zmW2iL4SIt@NHg8zcKk>y`#whlcOI%{ZVuJZZ^FnUclXbIoI*dSG=&n2eHt{N`!tN)7wppT1fwcOxAU42nX}_YPxAmOO)sUk5KH8U`N>+*FxSU( z;XzP4lYRv1&SS}T>%n6&0(E?d#$PR7P01idAIZ3O5vlnIJK1dN?W!rovaTVX_~1z9 zpn7w>5kRg`)Z1>zSklEv>_pC!225G;wqFL#|DMY-01pjC0V0-doJWLi%ImDBb41y@ zJ#QFv$}BuK*F^<-z9XzNG((;JizH5mwyA*v1AuK}Z7r5F?ufc;qjy4WWVqw4UL%D+ zsi+F2soeo8FHM*ZILE}kU3ZJA7{&Ik0L=m^kgvM%SJ|A0NC-Vvj#*-&DG}SfXPP=h zHm60*qdc@=_z<@KfzsB)relScdglz~0;$!h-%JniDxw2K8sFx6awFYOrv#*`d(V+8 zFO^uyKAqQ==@q?uDk;iLMxMRnZNOgHvmQy@N-El#zx)6L+S*i9UaA!5!qzmIZrqu4 zH2j;{YTIF}*z)dFgPYmuucmKko4)4CL>TFOYmJo8C2Cy~@H)UVCZi-T454D;&s{5C zcjEQDGV(wXkTlAvhw-F;TxTo6_4~-1Rrt@Yli?SoXbTyjw~d~&yA+ZX5oW5VM&x5H zCUSs4Veh4P{4q)%Nuu(nqc=zK+wm)9N5R2v0m@LSkg?5{gXw$n=PKV@vJUMEnLg(l z!mw4ige|k{@xNPt>CTGB3$}VYN9K7|ba~%#p+aBF^#_#bcGY}5?(&MLiV4Zfao=PN z9OVQ|xR4JQN4bc5y|}L=Y|Dkuv~u{B`5F^}*E(qKZucA5QYS*0(CJ_UzJPujRy55=GuDLE zCUFPR{2H2?$}g*#Px3Xu1$-+hd1qcSO0uEu2FpC7k=ee9l{SeM)l>h1x4Rv<`@tSb z#wiT0APg{@4;L^7%^~AZI$rD@qo_H2OK+WL9?O{Xp(e3$F?{!jX^WKmv=_Ne%!wt} zEks-U!i`tv1e1?!+*x_oFKAP=UiJJ6%!l0r;k7Wuc!srKJV{iSn!wS69Ez}1@(RA$ zR+Podegt~9(Ca?>I4-(MGllDm!rx&;MgXQ0YFU}7u5oVmkh67gkKB=zQSaxii-}vJ zf^Y0+HB33@$T0m9aqx%jrC5D3XEWn=*`AD&@f(vv(tQ;g+z#pcnY?!c@_h#OWegnq zP|!0GGL#k?T5x`GL?p$5f)N5SGxKS(X&28Jn|>LHA}>w4^-Cn~3^X;f*bsO7lghORS=s^|+ibJ|AdQ!hY?t#n4M2KcbV_*NL#-U-~uJ2VMg zEfo;f)+Q+R#P~w%pc~%6?`gybXJWE~7wg&0u!uT6BExG`uRuo6WSd`4yxM`{O4 zt(!9~BNs5!npIfqLx5`?XgZiV<^-wx(oXE6*S~?EPVL7&H~)cBb2xU%CAWc$*k=>1yvBr^@4igE!NUj*mMHj3WUJ#TY80S=*pA z6gy8bml7&SlIM>b?Gv{R625gNxO8a9+O5zsC@A@Gd8gTBMX(C>J_wBr$Vu0*^c7)hv#>8 z@2t6!mUp7nPu?j+3bsi2=yL#QTG+1c1A?pg=U*05;+X-hLvqZ{$Wr97;w*J_%z9uC zG9>z;oP@wqFk&uFuBRYFBzkzFBCDV2dX30P_wH7287Zjkn|WNmN~g{EE&3HFja z5TE!)Y+q@{5i5ZVbfn@=sFp?onk`Kj;9NI^upMxul@Wf(@Wq{mPL8}2?gOSm0BQ6a zR(IK{$cc#uDy(#o0wbg0)9|%v79Mc3ObH3 z?1yNtQa=}l3wx$y7_d*FL8khV4>+}llTI#u@x_y_vb~x(Ta%&d_^^Bo)my464b4&wJcYh9b~>8V)YAGo+@T561X4GnUE0n*4=%#Ti99GYE>_rh~S=H>|} zbkQ3iG#x`G2thuUF`*AVcd<)Y~bg;IVheQp)rLk{8nD@x(yGvy#b5 zqp~=a-m=FF3J!J+RrhxJJ$Z1GI{{NiCCmd}CNF8k3ExD89KhWS%n3(gmIb=3@J$~% z&8v=!{Va7`=%6kM0{Pf8P7CKQ`ysfpRxOq-;y2h@*3JQ2b!;5UQhU;%Wp{Pph$<{QstP(C3z&Z$|z@Givv`|p4EBR}Xs zWKOwe^&DVoZ-|{nS$--bRRseWDuA8-ZXPuMy&yFOG zSm2w_azJ^L(E~~j;KRWqUz>qLC&9mDUZpY;@lE>dtm0+jAl{R`^sxbG;tXL>=Q=)> zO`9)H0*o2(VqHDGgu$Yb^t+uH#+6 zu9^0G#!g?g93b;K+ll=mlt<$k|$pxsH)9Map>s zZr{Ra`n!=bNjMjl-b@uiYKEFTaA@2^d_0jk7Bj5`{K{tn*}l5Y)1H=zb_wmrYLZ`} zy-B{h&`#1b{pqWSCX>CM>^tI|4JSvo|Ad`!c6`ydD^TMF=r51KdjSU9BOjx4u|bHY z=Hfj$8ImgwQ?Bwo6(%}Oiy3LwC6i8`@fGN?zJ}sok#qr(h_3*4&}yOdkHC9kb{%UE zV+HVKr5fSaoiu<7vj9txG2eJPGl$3+d-Z-Yq-H`Kn=oZl-ixz&Og#?*TW__*7o}i( z{AEv|3vaXavTMW!wg@qA!L78EWy5FNm@19=SuF>arU7ZC~vw-waF~7!f82pA;UIJ`CML6X2 z*hGrYozm6neRd?YPE>;ND5&;o2YjUWJ^kOGB3OMxvltok(LZ4Xd&|Gwd%I1=dhPJV zv7BmviR5Z1{y{&uC9kzTX7i!DS`=*NX+!%iJ{a5-#+@sjE1bgS*MST3 z|F*A&+_H?R$qg_dG6rYe>=b7Q64;jIH7$M%OfdBAq#p+A7@90EltsI2O<3lE-?p?} z&iK*XEwevtg`PxvCS*OfUuUb(W#U4QXH4vUYO0|jY?yICRS0IK{~5C~#m6C}6g+xQ zzHh0Ge|#-KxAk#%bmL{6=*CH5ex-=>wp&a*fVw7_X674MX(90ayop}y0f_SJ<8FUU z{s#?|o{xc*_Bvc=?aoIAw^}zmkqC#m&s{CiT1zAkc7Qqs((xw_`0-2je|wUAEzNt< zZ|`Wm!S~zb$STH$n~sihM>i^^?zR@YF8`YIT=Z>^;*0(8^=%z+zG4HVD6U*VJqg4l zfLEV@-FsZn@01x^=WM+rd^C{j6?Yd5HqN_^dlC~vnp2yg zE)0Twxaa!2;9KAD8_K3@Ht} zAIjZm?c4=RaZ3`~nVn6~#fJ~pF10!o@}Ie-M6feM+<9^~`*Hiy$qe>;%el=FM(5SD z=RqXpyax}}3Uq*x;K|0yNlW16Q6h|t{`9*@;g@5531h$i2mys(p4Mcu@XTD{lt1}a z<}uqdEYhujAsk4$2QaIEN=~;VqFIX>91a@l;@NthJ4xcQpcTfr-)W@p2+1}876^S|3Z+^(#6{o_xVj(rzT&)c4c#bK`N zc36A1)ub){G)$!10*1TF`olrI%L`6XP_T@LN z&--aiD`(1J-rPCi-FF49dLf0o#OmVP=ev&euRrmS_tBbUTG$Hdw(n=Q^jfe z_+=jXhh{%v+MK(s9g_jnLMmcy4{fFhnwA$Ku#MAz(TIq8w`_;RAv&x2BL1oaCcI2`xcjbZC+^WBw zxpwO{{tfPSq_gc6$xD$ENql$bx5sg>_D|GUTr|4OXQFLbqSZE7$bK8Ez0PzubM;P) zum`t+3|)mMAUgsks(%jwI<=}V;R-4G=BbwVZdg%-_wU9U@`%M(Ce^uLP5luciweY| zrwFj&w3Pj(O4$+rIwl7bJBweQkYX?6IcWQ)t#fCaDSR1l^XlTJwC`mbv-(AU4x^4F`D?jb7S>vQO*8;3-eX_A#UwLA+UtpNU_g%M| zG7_|3Z`0a`s~)~#8G{X|sC$@y%lrw0d%b!OF@Y-T<{oEr|J|J7Jm%n~>8H8vJA=C^ zf_|NE9`T}Pk(Wi+oEI6nj-+L6&-R;cyGOs?9}QoUmWGW+x4lv^fK5?pon&{|OxNL) zQXI1Yc*Eq~pRm-liCTl%B=;a;a62z!%B`-zIUB%f{PmXq`+ov(HC1rJorIFc3AJW& zADj`N0^~2k=;#4+gU(a%-234>S^Lk@uSD(Hjagp*uEQRIo)eQ^g<4H?axNM2yPpv< zHCJ9|KrxHu`zoLGMi_3N3NPl38};o8|7L!dOGz5berwLB9hoh}Op7w?FojJLX=7yf zO#Uh_zK;2;oHs<3cDv7R_A1%*8zwJ~kXndwZCqb+v!83SqP`)*K7d88tkQGEpDNup zX>j#R^4Ive2Rva#8xEbkK-Nh*;Fwb6?`q`EJnoT}K=|B-zAEr~_hP0tpIG`0titpC z2V|7EpZ)!$=%K)n0xiGfCQp_UzE`K{!~BWU%GsVz__`do|8e^~da12RIwEE-kAG6M zF#L_7%uRX4d?ke4Mp)1k)03tnC08!`44`psyDnai0|$3%YAcmuIFUUcfk;G)MnAgP z74C7RbMe&Y3){2^)4NtzM%@&Iv7;(!RW_@4h241~&d`^%gyi4wUeclL=N~hKY3tvU znRD;V*hV??y4}2uJ+pf^*K5ACi2awi`u}uZ6hO9sq3JLEK(ID}EU zZWz}{*$7rhgi$&4f+^FDxDuuRs1p7Bo55!;K#_~_&v}V3+sgtG(L$N#xp0~^CzzF; zhr*AlYArESLOhYfmqpf?nYBZjtDpEMlY;oOpq7tZbK{FTDP7Ed8@-g8oeuZTSMuL( zs1`Y>>K=ImW{iNGfDR6clndJF9{F)j9+d>x2};Fw!$D_Lm(`cDM_ImCTvy^uUc9d+ z{xK&H*Q_h(a(&hWvHAAY#jgDn!Jjat>gqQtW1PL_2)y)IVL%7$*E$h?SZlU?YY_b@ z!2fA31zlb{JC&9cXh=^@+5bg<#@#Sgj4PJ#nY=q-o6M_D8EKX5(X~qqulQVn(P+4e z65o}m%WdE#+zDGc+R3^`u=&)Ab5kh}W@&(qR%#cXaZDeiFeNFs_wPUlD6vP?RGSC9 z+T;%&wi>?jT1Lq|9tp0NIbxhkc;+N*vIk7H=f1E3{-H-7Iqi4om-u8j?`0?pds<;)SwJ!CMG8c*w443AsCHCXs}Lpj+xU99 zJ`{o2?TA|~7p*_`jGduH85Y)IKhQ~9c(t-Y3Pi-YUHV=sZUDYXJ@xskD9GoOl?kMz zfF8p!?$nt5B4i$I5tSr&=}Y1EZEtgj7i3^75ngM@v~b<8SyOPh5T9wYHB8`-f;$QP1TzjCf|I&LV@%xLk;~E2X9+NzfbjNe_70ABzYQ$OhH__02hpPnhwBfj&r43|vWRmDT*Dmj0(IRQHKuSI8& zX3x5r!gudImXir(dGgQ=3|3^BSPttaY~cg`$3DGjJRgZp2f4u$BR@`Tn5p8ezw2aJ zBFOx%JsUx@-e()%pe8(C#M#0p)QANLC;pJkXo?Oc z%tz!M{GQCb+p82i+dD6L%2n3|{MDkxnJR)G1r_$G2lz>of}rlS9S z6*|5F)S&uLpnu?kG+^f@{Aswrqw4U>OvW{0oz-_PwS+WMKiJ++f!ki!aGu{w(n24cjZlCo;t=x#nyV0zZr`lC zSvvCe?tT7ILJ9oJz_jHuzlp`kVsJ3p!`YV!KVha$uoT1l8@044T&bdWrGA;=7>>G& zrhOLewL2FJuAu)4Z`X9m z?b5)S$KBO3U8HR+Ep@-;B}@AwW}H04%2xy@qTF(648&x~PgtIE?goc;(Yj}6R6QK0 zjpKrp14Ovh3}lc(3nn|>t}$r|bPoCgdV1<7%tGU2){G`Q2pk!ft)_f`!yEr`o&EUv zZRlQeFx~lJx+No2QoPI$9eNhw&U4OU^s1In%2gR2_h6nN37Pxg@VbB%eIqb0^%pNB z{~KJdW~0Z?ANbqK($)W4UN6BVzs7T>cSdtUL2`~4XLwW)Q{wf) zjdx}#@?4I>wY=JLUr9XsuO?DA&?GcF2*VcLxRKLc0n)x4$HD7K+$Ra5)D}=28tuot^sqZf9Li7q9Z797q+S%e3HT=V|4&Nj*QntbiIw!yWg#CnVu?KZm&lu9U+)r3K;K{OXnJU69 zLR(qNIkKQ$!!50> zD86=;B3Ap#nen826x*NTeZtR}I8GKJC$}(fFizF5uqFG3hTky1?r76i!bRv^}>vVa~_(JYm3b0GCA265^SAxn~=tDYPV7?gm$rrTdVtFu*Z+VJi3bh2rZ*5&ru z@}v}Di<7svY%ok`yQ|zSG&z@~D43S52)Dar#>0KIPzvw!;&UG$z`RF3f4I$T^%)7f z%ZffTt-|-j!ggEzfW97x(Hhn2&djWk*F9JCg9oz%5?~d78@+w=KOWabuoUaT{TT-m z)Ycsa7PGe``zi*%7Jhy}o*vM;p~`JQxulhIlS_9^Y6Dk1Zk*!YnBu^xxxnc9VLZpD zrX;T@z!Mo~=Q@+M(TBME;+sSJK@<~)gf`ZOjLv)+ln@OaQ+ww_+;vJJ-Ayj3y_{F3 zh@F|g`mo9{m02rSOE(yk;PUxHW$B?ev*cuxRjT;6uumTL#jj6S9~F2j-v7!e&cdTcJ^sc z*xG7I*Z-J+mz}pH_i*M5zf^Y`S6yl0S!*nIJpUBOqvS6`egMjHR*}ELw+;Ix`OD&% zJncctPuNl^dN%&o7Hjm6cOm<=)S`HHCsKE4l=@V)+f-Y{50<;B9^my)8g?_;W)@m= zLGKB!v39$b%v>G?jorOxdMdp(_<|h|Wvx?2;H+K#U3A}G!%Vk6ug2kOna1G-J{}QM zy|s^*gTD@rhl=hdVy6SX6%IXD)OQOGt8O*O6?FaYc}9f*1X}GQE|ICgGxD12vc0d5 zOM{c|n?r}c^L`CS-^VMk;mFP6v zKSr-VG?zEA+QIv!{&GOh$KtJjw#C3)*4mnmT-`pU$p%bAK!ZGX(qg8&y43%hMfa8t zT<>7t;jI60p!`3-#j2o=o@kxr{WX~2i_;0PJLHByH|^7{fUW1U2kRB0u*{vFDh%g| zGHu*>P$EEWrOi_3P{=!Jq?C#QQZaWpil6uX=A{nX%zgm48gHfOThyw&?Fgunr^~$< zhzG;ET>(VR@(v0A8kRZb`^DL+zgv@QSbIu*21!at4+jGHM#|AkYX#*xFWQlw05vg~ zmToiW1o^7+++kaX^dzBS824HSzDREUmp!umsMB9=-bq@vcSF=k%UxM`{p*&n=13wm zetTW^cC^)BK{f-9a0{5tKT)K#`fq_^J1+U;8o$qN0ZOXea71<7>2G#AkvZDCEe+d# zao`dG2TrTztb)$0gyHh+%7x$ws=*ZHDxhoVDi>}&@k)E0&gMDfofu1GQk1rFlu-h( zD7~Tb85`}EPgBq_KCwr83F;7_T=M^+%QaJe%z1Fo7d7DA{u-C%b>-@f3pZ>U+~dV? zX)0|2Dvwx?{Ij#SsREFu$rGfN0{GE;zp~2BU3~TymjOxjF66_B!Yr%d0n}WM|Ih#V z&;0x+Hva#xE-HHe!k7B@SWn5MxXYVC7HSY|fo5(qI49T&lF8GZs{T1TQzv590+L}{ z{t0SC{%x;=5nGao+T!E{toVMbsV4v0*D z)#>p5@fvu2uE1~HjuaVba}kj?(-f`LGmXoZD%#P}z3$wTT)*|r=Kn(+ANKE1k}TJ_ z1qgL5cN6y@`M>fkCq`KBHB~cs4I2(_{Gki8R@sE8P7R+Ezj0<-*IjsvnkvKCJt#BbKc~&#Hp4%;?N5kwYsH}HT6ONcllOPPZfoP(Y{j+@+sLSu znR2Nv!qQt{zfoq_Bntjd-N82%UD3a-_9j(yF?+j~#|x1TEoE1udo9?B^p=xr)hnH@ zp4Gp-Q#I)kyU1Vo!LG9~_c3m@$ef}A*vYj5_jh(!${K_HeH^-6YiWA!Tl3Tt`A0qL z`~J`Gy60JQ#A~zH_}4h;&DIhQh)|tJaIn_PfcMi{oc|jkAXnhOun`FVJq~QgllDjX z-)DdW$G1{QfXxs{Wecy=?OhW9hhKmUA_09rYHareF~l;QfbR+2V>F-SNis84%=#}i zMfP4upWexbg%4Y@#5Thdj4Wet)(kY6R1rLFjiA`_VyU*z`;c6vUXCkEjnr7j@ZG@t zro!N2ps#}q6gC1!U|%=@7lz^io6wQkPpU89&pdWI>5k6$JGOrhT%JRwA{J)X;q2Hq z>QtB*YpAr4)BdZS*xkA;I)^J`M5@_ugU^r#6SN{O4wAgsMDsE%!NiJ|P zFZ3Shxc&(0JkD@Tnl1k#N6x%W94V<6wBBQiFaxq10&JimZ+#5(GkGEgEL`EpRDBdy(IHTqthrBOm%wxGJ z8%I`2o4jkn>YrRM!aiSedJcsT0ui>DI>86w{h*T_c-3@*?p|ro>@N?5FD-%i7eFDR zLkKYq^5;n+KWvkn;FL0{d*&2zoGlPlllhTs+uR@tVHV@nDj9vmqREhf?mL^_gv%%QFLm(*g%i#%M=Fp6z~dth30C|!0EA57yP0djDbLH1u+r7 z3D|@Xu4BrBQ@%mz!{iD`5IB_%+$I$aK@QIB96dygt2`8$XQt7T#5dwd?t}8_a9-9y zPFr8TYGwrl_6uiXoY?0UzMB9??7 za7H9X-L3+&5zjOp=48=GQLc(Wan+Z5@g{>vE{FU8SQvjrN)iq8)kJN=`p{zIkX8)O zqkWIJyzJsAeujFQhJ<~!jR}xt`RgBx)z$(=2xuF+Npm20gc|u5%e)UH1(XsnI7(v* zHe)0|)EUw>6OH^R_f$p-$|6z9aQ@aJtYd2ff;mQD89pM23=Et+@2X;>>5xEcK2=FP z*tajUk{F=yk`KMQbt+tl$@jf#^(srJKU@f>)E`6@;&eX$W%WaG32stxf*MCkeGWs+ zRUMXb57KZ1A$~r%Y7a`3(Q9Oxg!Bk1hQqy9!t0;70=K3Q747$SG_~6#qUjwkN_>b6 zTg*pAUmpk3hr;)$oty!V%=^OLCu0p6HjHoSeBbK>3dSY1PrX(d7$5NMc=*uvOrhy( z56?aKL$}mv)%Si~C4N#L284D^7f3&tP5>Ow7K9WM!N9kKZ~ibnQ`(>y-kT!>>HE+~ z;uDUt@asX6C~nYO5Vb%Rog@_d;ThW40j>&FrT8R1-{*&6GLqFa}u9M)zcRk4387{%v z$|f#!CL68KDS8Z9H_J*^qhmR$pxZyZs0diU3>o&bw1BGa_;FW7X5R~%^{mL0?7M<{ zsaF`E}d z`YYb;!Cw#O&7e4j99xS1QI?cCb;$+zfH0`38S93yNaa#j)}a_6{P3y)$gPyRb1f%Z z1iET48j6O?#C#Mb`YzTPhd%NO_Vox28qt)-TDNlKay<532F8$e0;N!TlcI9Da__GD zxjvcqqF-X)e|?Z*fq7czxypw-$Y}k}T-9gc-s2pAPp#(aepjE{vTDr)VU{&yT@1~= z;UG3){hEe0y)V}koFIw?p{1c!O#j6?8!x28NaseZ)k5dBm219vIIE{!dgwjUFqB+g zLo7S@m!muZ{k)@K1b8_SqI_}N^%FNU#i5XSD<(0bLCGuRR_isrRiYra_8@_b7e?CC zsBqxXC;%nA50h(Al&k3-?q``4^G`1%?|D5IVLgH1LZ2ncgm)oe95m|+V=09r30>|S z_U(^+j31F>l3Fycx$@L8HTn}P22FM`XLoUw!1!=ELOr@u>f^ft6@p!R4kyytzM43x z;niu&7Nc_oMNTjv-x{94x(L|umw>7jP#qp7P!KRL;MRdMHyRbyz(uD(YP!HZ5n z>eKN9!s_j|8zBph!bHRNUesn(bTRfu2SyOWP?j1tJqMzG3($Cu zYFv5ybRtKcQU9&LkA@I-{p{>KW#n-)W48bd6`I!&4QZ}w;4ZRr12K|7046r3TyRzx zD+i=D6KS?0Ktsf~8Ge0WQzeOR`9ZwNRkmk=kl?D#ICrlMDk9tdU+lekR8!fy?i~b; ziVzwB0U;(`RQf4*F5Ha&u`Xh z#*U!zSvR5|#G$w2*6~ z<3$(tw3dc1Z9m3o8yz0Lmn|h!ah3z9+#1>%r3douff5A;=0h@RXQI<6qcWK9?+*w> z)IbDF4tvmlT4@Dk)Q!#!r5nto`m`?tFF;alG(GWDcq!e2>fK(oPy)j5@YgxZp-mn6 zIicrYp$Pa#?5a{6vWPhHTIvxjMKelSSDM}yQCGo%x5Z?WkWQAdT7SUBh{6qkLIF7+ zfGVoyJCPHd?P0S$u-%n~%sOlH3WUp$LHHe$H}VbI}owJd{3t3B&560Sa2g6A`GD6u?% zeO{kWNsD~lTvTR*x%G(Mj+;WSSx5*0g-{_#VbW*hDjnO(FtMakHj`sn48ta)$Vem6|!~xD}eP(?kCT{tT2bD>6i;|0$XY~aF z_tKRxWkkCpNf@I%_5_6Hkz~ZK3@^y4w^+=!QidbilJd?`_b9+pKW`CMkx1PNUtHG| zfkpRFh|)z_Qm&qITXu|Rb+$9tAFD$l`Edb`J3Pby-FI{Lu60xYou8QXBwT0PY%REt zkNvhJ+ir3IJIM@*O!`v}?!)Cf>Vde_j~Uo6)3M9aSYL4ij0dYy6ypu%bw@M~<5AcD z5VCvi%hQzo)AoG^g_b{1MKBlBNAlIF*Pk@}sSbbaM-A;bef_A;H`e(Bd(YUXnH*fK zM7{m^)h9@kP6YMWn)=O0NUkeyuL$}H2+!}hJ9_9VFTZ_*MqgQPb4b&~b)D(|t0UTp zKP9kB*Zl3FBg)8o!c1KW&W$;;g3!``K|ytC`giK~gg>O}P8tXs4>TSaK`J9bCg-Vp zxnsNAZ~tkN*rsXL?Q=}iU2dBCPkxr1Ka~}lrnp)Q{!nULs;KnwpUEoynQ59otH$wj zwa$$gplf*gjNXq7`W0y!M(2K)CUtMhcix)^f2glDu_n?MI=HMlW20@no$x8e^E!6H z$MLmms^6nfQ~p|MA8D>h!E;5EstR@Bu`5{_1xyR!gt3!1ga7>gG%8b20nAaL-3jph z4YhH@hCe+qyZ*iu<==RcQGaRj&=k{Us7XHkKS&Un@~6wg$?rR-iI*k;%+wG~GBZ@b zpVIf;?ID_^Eq33@(NP9}>GlX+^LGVSD5LKzDt~SzkABg_j4{%Hi2-1-%{;MzmT66A1e4u>&JnQqKwfOA^nW9;<4jsDj#%|3T zAe=b8<^?=WzFGDDcIu-&SP%8N+gOj{*16*XYdeko?f*b;QH_3x2-KU3ywo6`X^hjs zJ&m`bedwJJk(0Q=qe5NKP#7+PUx_h1AX))d1Nxz&i?Ys<^)gth3KL&Qj-Yqps$HTn zv1GjN31`x@G5ZlmWOQTsOIV$+4X{Cx?|PnIe$-9z(Jfre#sNYaMi~s07ae)Op-#~! zNnxat$+Vl0r))d6#D$L)l7QK)q@o%eh9PBuDn)5bmkOq;o3XmT=W1VJ0igcqG>W|+ zR13=TzF`h$VaFKzvOSbht2Tf-W2x)u(pCv3-ldYArVP&+%AVD-Y}~=g&{4^z5MoRLTWbwPP47I`a z_a3`D8(acSp}VoQE?+3e8@IHU6&ABYJN@pzdIj0bmB?WQzCYbB_zLLc3Cu+zr9B9g zWlNdEtxrIc=qa-xfP&`Jt<*Cq{$fG$NsfJY2v*14EsuJrx%3m1&XieA^T6&v{4Y++ zIg+WT>QSdg(SlVLNTQhh!B|SF6U5WCsR2@4*A4M~)jj&LqM%LiXdMYd$&usR?jy$V zD+!??YXg87M>0(inViy8Wd)f@Mw7s?&8pdn%?Px61_CE?4h(&vWYlJ3+ma8TIy!CbFu0U_qb?w0MadIAuS6)pg0C({x%FZ7byv<41Y-{qZE7j zD`mb+<{4L}cecK9`%4f{0KN)tiMm8WN;0DEY>8-0jiGww4jF{BT%x#K1@`#jLQpn~ z`*@^yjvj(}>KBUFYX1!tG(xis`wf*WoQG`43l%+2dnHkkcfW!NweKv&^~nYO(Ts+R zBNmDS04{|!RQ>=G2pbY2@5dZ|;_;MuEXb}F;OIO@)$AEQGmFRtjcH)8#Rm}JT=&ZuLjRr=OoO{p%bpXkSn>}1RuR$8z2Gh|#sN);7P zq*kHDqvn*0b=h-~*APZ}Q`Im$E@5i}*tzL> zU2c^uG@}k1M>qTgo5w{IZ5euW1UWIOlWE7wdD1-ejwXB*ux4PA zF=(9+{Ng`$IDJzuz3pxe_epNDQt`_|33J?6u}UUJwy<;~-MoDdAP1mUNqUh4n;^^; zO{IFF`IwnIXyya7MXjfXBiLvpXTwWQ``AUjqg|o953Uz;p}q8Ykq#l4l*pN?N~g^j z$khA9txdJ9Q#-*-{#0nZK10Sb93)s(1w_G?646!n!-b(c8=izVgUoE*GZ14#X=&B9 z2Q5{*Vq^vLD@v={o#qIl6b1u;%nN#O9~AOy&;;KNdx?&z+P92rR2E|fvMf`h6iboO zx^!nDvBv+UBP)&!1`#HMRC?2x$QZe@4`5KjVJ~?=Yb4gvrA}{~TbRPX$F0`P^pw_3 z7^|LnCYRDFcdAm|07q=~B52Hy8%H=kia7`{6t&$=pAnM8V}qP(5zmb4>(4nnt1F9RG%8cX8=urk(yTS3d z^M-ay>_rCZH=t8RPN<3zW8>*2Wz4Quo&bP>2 z$Q<``=*tn|{&kt3n5nrf?p{$34pOcNYArg`#>i-j_NETy*5#>(>4gQ)Av)t~Ce^{F z^~a4ZTgTB?2q#smYe~95I1oCn}p-kNC4%h+USl&fy_|gWiX{ zdZ0Z2N|jDU!>^=-5yuLamU6X}I8_P0gy9L0B?rhc!3k7jF(?UGMzm=I^bdH0R=hT!$? zy6K{dS7~3~P(7ISl^gMH>AEWGs0Wp0)|M89br84?oKCjO`+AMm{O>++} z1PL8G3T1lHg4%S8^5=4PREU_agqn*?!qx1AhTiT0}! z)hj9Ir0NW5caNd_yrUE8fp{+TCB!ArjJdajtiS1^#~>0LcqsfROzvjLf5Dc4D#Sa` z2~h=`#uZV2fQ?VkEeI#v=xMHI@lRQXvcA%xBvx{Pmsn*f3AdtpSUs|M-I*LY za<(bNFf_WdJa$$FH}5p%hFqspjR83{bQrx+bu)rH3t0_j&c?WU^-|dL@)#>VlUVkj@-($>UmA$|Dn9^9=h~|?_kTj7*zo;;PqNuJg6UGS`mLMm~`~9_tq4_)8o~nq_ z81NupD^NzZb0RT<{-Uhh(8{@)2ii(K0^O3k8>7b6I<&S976y-{Vv(3&C}$^3s>8)%%d&4Oys0zc zEEY3-EWoI>-nFcu#@mgLp9PC9J*_f+N=O~U8XQp0f&O9PIO&<9L_g}csQXwM_pamR z!`w-$z7c1pl#JLNzX9X?p_}EG6*EK>7iq;vh20g-f-nnkMu^?j3dVD+eX1lSmcV!L zl%Nxp8!dMtXQ1>(dFGuu7py+@f>_aGxv6#BW4XW^s%KlX=oj-vblhw$MTJZxRP2L% z>+MNh7k$8I8 zpR<BY5$xWwvMT0EZ-=;_@8b`! zTLMV(E)c>6KOJfy84!Z8iP3h6v_E2>wk6Z=x-1dnS1^cN11;zsV( z4sJ2Tg4$YE); z$IT4sO#cyeAsG~xUS-c=U*0?a&4`XLIG$#zZWBFk0Jm{E;kJ4MuG#IC{LQ2Q&9C~{8d;>n5l z`N)sG^V=i+H*=oB=WhoSe~q{m*+q5?3qSkyD@gQN3{ckp z&X4`z%x^T#z5gAf%^(16P4P8P{=Av-)nLamx9Iz;nW4&4eds^}J$G=-ctZ zT=mQJwYy{0$5gACS#5pX?Li#MlMS{}&8$V+znx7vKKE{^JF{1&jhg#o)Eu2)l+J_j z9Gp3@oy|Wi$g;*v%Bt5a2)g#;bi4W6i2(96A3&aPK=d`ZR+C=%I|=E%nSPX?x_D(W zUBxfcQg1PPIcGOM#t?>6^w$pg1R0;YgR!R8+~qgey-mKRkX%lgWfY~;hn`ZNJ-kdZ zWyu`O>d$)&&=c&|h@^i+eAJIVdd^29y$pq#Pj)3_y91Z&7~f*UDwQw~9A zr!O=PH>bS$I&Z*OpV|-6JFT^9Z0Dw2TEF(+(6OTy@ByB<1*ZUKNhK}ks`tk_Mha%R zySA~>wB}1&ZcRDUCJt_jDEJp7(D$aSbRKvH_x)o2%l5O>TQ9MkIdiIO?0(#F`t+1V zUvs|61R2XQKLLB^+38=OZ(URKrXGW~xuQ``to`3~Pp3~G zcO%!Bk6A5F@vc1S9`$wFIpevkMaw-$Vqfz!d7AlUV1+Eee zxNS+jgsn7x^L6XVv|Gn&KKA7u_{qEUm<-tIa<{Mjngd2j2PtTsl(Va}bUylIedyiK zTDs0tcsjQ({g(}E_;HlM3nsFMM3dthv5f!smj1Wg(*IylYa#{zs97nl$uVOy<)8Pd z^z4Psr$;EY>aSWk$89Ircdspt6)0@A(?Auq@(=9je}vdP{(m1+P2pvz|CxKbW@GR6 z{gmo2vg(y9ZNlg1Bwn(!+A%E#z3sbq&m+mTt1Go;U7ap7pEei!)7jOF*LZARhrefm zveo$FDV37S_kWT(sBe@ee9%(dzsqz?e=`t&(0WJozFQiZiA8S3RmJ$dmzbrJt;r6z z;?E?`gxh|#Z`hu_c*-yTm2cDko}KsqSGV*(s0#iaf05Sr$Fb{Z}EtOh#Df z1)gA7*Wk%_J%0Y#$z!ANUE>t7tUs(yrT`4ozm5(;Mf;`z?3G^|cLWT{y+*=@>1o!w zsyn5k58XgAE_;L5rO+jwi5#~xX`kj*|RPiFf(E9|l{>KIQZL zv-Z?E`{jbLCe!Pw5;rfd(f+OvNSGhw!Z`k_RV4doxMn4-S-n=3p9|$z6T94`3nzG3 zGlcX5Q1XRIi8Ul$ka05)y0f`+EWia!jCf34>Ly4X7bF$Tpe}&u{35S9`~r|hnnqg? z!hABcg3>M3d)DCu&MD@ha=+jNfE6c*3$nElZHO*swR5Jc1=$wlS;z&+wA5S*;ouee zbP7sp2%c`v9egT}Ki#HSm zs)6bjHB2#^)`6WvTTEd|v#-Gf2{9F~-LJNM1+jQQ0%_sPIH|TdNZnFVR$2xmV`G*sqw%;# zSPrykG+U$sL?{?zuc%i*J03W)W+K{rzutri(1-2gDb|dEsaAp^WhMzUGAo$8%2Qb~ z9Jc3#dTv>zSEC~N5$7VrtOMcTB}CVr`%c_lT8xQL+*3dSM1rPI*-)Y)FonQRT@- z91(JDrAHn&)v@5D`yShx$6X<9X7bq(wr*5d!6Ru`ErC6OmYm^4{P0!GZNO@a2Do~q zbPRO^^6G*VdxfO0B;_KWKLF(zSGZDXdl8m3>_pN6$o%N-B&$;T^nVPNbCpJQaSlD-jq1YSRNz+ESU~!_URvi%qRr8)G zZJ|Z4r0fs*@&F=3Z40qu4?!lW4Of>fyCV)@#a8*z%>4STwOXqPy4{1BjM|=6VKK(* zZt#SJq{wi>0n|6LwDYpR1k8I9==vnJ(OIgVj9^6B|!gKngX?OA(kLCIa2P*x#yJOZ%vBK5rnE|*O z)}`lzm7aANjB-J%6X|lcC73v4x2m!tFjMnW#xa+FfM^1mVT*dj3Bs(g@E~ao5_2*9 z$zb+0=yRczZL+fSSVv@e8Mpmv=NVMAD2;GPpp870?Na4c1t>xgSp&9&n@F2O-GS6e zM}>qstoC?TkhHf)57Tm>?_1^heHUO>MV zUyLgH5R#dlfxm+*R<{r!uhA)>Yjp={q8TAyJ&^$5cVX0vtPo?_)VtP3pbcU+tP*3U zsDH7bkLczqcU%l{CoM0K#-OFkZ%dd-lbf#{E;xwn??M`DPxl*yMlugv>??#e~F3Vp^^W=p(>B zVv4L@uezt6H&ppBOKjWl;5!D6pHDjvh5BXlOK=VKs=>z=; zf?dbj*M_sZHYU)adO*K8pGbdJbv~28a6Cf}nd>a?Po8WL6A?S#Y5{{(!2pfo1Wzo^O+ zG(1rvimawrH$%47wX+t)5>VHdbShmQBim#??UCWjy%*v1k)+05P1M7YQ{P`BhZXt1 z?#a1$szgRYmK;%o@yE0y6$k+)RC59Ho82)_NffpZe5@~y z8%&x-UDB*v{g`r*Kf##~#|To&?73Q@LrQ_2OtsMLv_GS1Rh zE`mI(-ALA*Kknpqjwkh{LrUpaScjR}9TsCrHbnFw-j`CQ(MNe60Tv>ou zt=xcK8(>!m5?{fvcj|g1ECdF+@hSM(6s~uxyDn^BPg@Pw=SZ+0n3D1$@P%LI8Xtj9 z-IkCOT{69&khIClwUDtd9n}?*X_q#})H{Hz5~&UAl0424y`qF$33%Ydu#;l>7+s)( z2`dUM3FBBH#JnMn2@)Vtgk!fQ$u_krgJs?zK>Ph|@|nDW@=4As=Rz6QXi)eQ@kvi; z=d-ue%JZyc77?cDJ1^0mQje)F6ON0Fz^iwC)q2HJr4dpl6)`L}v&4?@rusw~ZVD3U zy-mh<{(ajw2ZK0}H@o%`g)L7geGc?D>_lOFag)@c^ zbCKVU!3~mJ78>t>C?wQE*}@w&*7GWxK2U#vRx6$0N6tf*Ok|64=Hx_@FE2CEXo}oN zy|D91HevZ7&$n$^enn-**hGUG1zxl_L(H%x3_jaZX__ph*Rv9PbrF`>!xDtugP=eH zqt^uy;K1{CWfyfRRMUZ9NwMX#v~W(qvNfyHjk1{H)}VGSgGot_0a$N(F>+%EW8W(; z@LSq^YZR&QI8^xGEpo6tDc7lHp8Z6;2Gf7rGs_t9{7x6PlD<*jtDz$$m~h_ z$TqYZY7=zfQHq>%W zshi7GIfO8RHRYA0iVOXy^k^BWnjyuSF_W#ekc&=FUp>vQ9|@J!RHAu9T;rZD#NQtp zz2tNUd`uz5+rQ2g%X=j)H&vS`4yoUO9_L+0R=j#85)MC<$;ja16S9{eJ&@Yr;$9?- zQ2RP;8!W2E&Iti6Urv@?{70CTCar41-Gj~Isp};%F_DcIUweAD3q^aIB0RR$etusP zfia^A)z+%pRZAY9=$#G^s^4KTG~2B*Izu$)6i4gCBqtF#1gtufs#cNhkSN9GP%egC zwLxq?PBNzud8b3w))15LZRyofHl$P1?pWCP>eZW5G3Vq@fZC%6o{+M5qj7Wtc1B=?Y8tm#!+3VVbl@lD`hdz5=Rf?za1WSo7qcP_}o{nED!c4o^i_AggP!OI(Y z(eJN;ri?XA#`uZyRVR9 zr=DRiWhH=yJt)b!Db{94M&XIxrSSVQ#j23LYL*cPO9`qG#?)U^5S81T)w3!7b%C?m z9oZ?_IE2YY$eUqeqT`HQ7@e%11^kqvq*AX}@AlMKj3)Th+-J!M@uKP6@GWC>CWz&v znNf9-uZ8?6xCb;v{NJ;?3QF_KumAscHRJfcFWjk?4+Y_rVH5DuL^yc zk3Dcxz9%H$f?|t zh8nUkBe0!n_&A@jtjb%;z6}1G(J{g_>h7jU^2{zIb}KWW0`g3DiqQv0hixrNf4YQg zFtBS#J-@Uia;^I=$(ynmA~F3}e^Ig8j(QAi4i640gAp>M)^n@l8qk*ZW2d?MO6^5g z@TMxpm|)i1`$(`n^>uR@HVtopWIz^Q| zOGaoho7Mj5%1WfOkh-aZvl!uwN(FJ@V6Qh4PIhmC39-vKK#L}B=f<5_9rdTU{=id=&p=&ZpY3&u!uXfu_ok&?&M zP#DLvXIl-jw2dYCb7Em2ueOQN*^h-iwSu2>Xyxx(pLnIhr;eu!F8gF|M_q0uGx@ z;~uJ86UVRtNQ3A`PTm{%By)}tVq9y1iwEkjn(Uddozjlx2tB}~29$aBO#1WJ}# zPFoGvd+sNl-Q(0M4&ulCC@bqFCj^ z@9;Z^oC}0T84&?)+i%nYDPVE&WgApJPHAN8$3bdjxU1O)$Ve{oP$tQ*=U(nLMs7=| zQ!zTO#4Cj0(KObShlnDV*5Y?lmBAdRV|oYhrcgaonq^-XXwpL6bn)SaqLAn{Emic# z1L`HMxIOUb?zCE~&1$QqzkUWyha~x+hqGG#v&X70^4hf@{oP)@cDC=?=79hF`5&H+ z@?C?jSI=L6Rh?1cH7A?c;^8lsAz#JbzVhZ9GBN`{KY!$mV!IO80;a^6L~C2vrU^g0 zpEe5#I+l-U|IiWP5iFj@gerJ%sOJcmsk+C+-gds~PUsycRx)yKe^re3t`7FYw%~7q z8f@`s1Hpni6UkC+#%Lj}8<0Is!M4S5R83%Z9Ynnnb@Hph$L+U&vmM-JY z><;tk$f)#aNOYeUCK;fvzh6^$K=h$!;QC;6%##c9&hi_=(&B@)qp44509R=$6kgV0F2i;4>RCpWz>% z!Lb1x{OvHIbIWh2wu@d*e7=jtxMcgEv%`$PvlXG<{Y@m;t`IK#38N|nK@QLN-TK3% z<)~$@>GLKYzxrs({-*E%(2RuDLh;Q>GEf1+n|(Dm*`1r}=vH7|Muv zXomaV>_3uh{``>Xm+UwERjQ39o>*?IY9nF%XM#%W#CO)Fb+hx+ZKICVZuyxlykIc+ z!)KdmJ5H~nm=8wpyRG3XGigxLKG3^6)k%IOf5Z}B!u*kK^Ve*&C*Ma`qdYZni!>!< zw0%cy`OXr!c*bj4}k#U|71$9Gtz_})dU@L`YS6S>R-Yf{b_hAI5DO4@9?LVZuv_W z<^Q)H(8sNsGLwe?a8DMU`3;3F$I8Kb`a$}og@1EjE1K1maAey48pHbOs}|eDp{~q$ zV>3Qm6E^PpU-~9doBuAC>o0S||EC_%e_G>9{T5Hb{pXCPy|@0AJm#NJ8hO%Ss%5aS zAz{VA&j8)KWwd#8l6<@&@Jz}?zdc{3X&$THjn_w;S=;o!5xs0(SvE7MkNf_dZC|-- zQ1(l)b}JPe#lEwIMTZc|NW&EpHk0(6I1(Nf3G?{Ycw`2cZG0H5rP&~EQsQMS#umBw zB#v((LdqNEYZ3M%p%MeLuTt$~ zqAqQIwibMWH)_2pJY156@rT)U4#wL$f;iVlrE`5VGx+aLc`p_0s6v4_ta?%fxt3|9 zZW5uZ4oJxs%FReDbfZQz^*2-q_iA)e1A}{vbctRvk*qY&qk*bfm$Q^&G)b62u>>cQ zItPCq-sSf$H|*S3@0!9{2(#nEop6R*rM7vS4rRXUD0oEMid@I&0~x+wITT-5G?Z-; zyXSQ>uTqfcsEeJ6#O`g*d2uabNEOC&B_&di<)5&~YJZdz)-%7ME_*@Qvlmeg_d3QG z%lu!`a|9N&pEK$8Tx3b`NL~zVEcb-|5vbQ>1Jn7W>4@iMh&}SV$kl|_DX-#|Zab=O*C`T5Y58EwCr&|)^r%>l;yR%{iLHJVcilcsF<$gBc9TXxMe`1Y*dQmNw`T9 ziBRVjWZx4629r_$8f$ChnuLx09V-zzIwQ<)A|aBUn0gO6w;qf)z}$1=mmy(uc}wXn zWUb@wt$s?xm( zYfJM}B7Z}r5-i6@XNH;n$a^^G>Zx4Dwhrx`(Ht3I_r9f!FdfmYEgA4N8t@9L_@Yii zsiq+Iign7l<*FM_z=BpNB`4FIluOzjM}?*$IRk1)qv-;=4`&s94iUrq3G!6h1ma3M zXp5A~E1O7LUOg-8pqfj&?MIC_Jorag_F@;jMW$DB2vNw<1)n-gsSBYQ_}N3}Io3#> zjAT9lm|&21Bi4fI2p>}kh=t4P{h?^a~Hqr=@y zvc-L6L73Y4C1|uqMY`=2#m(M{9z#n+$P&hIyV%2CBH{s>OR_u;E{}bnevd81RRf#} z^%U|BauZMR2^Jd)siCSI_jz#Nj`pVnG=Y@N_VUJ!gb z^YayX#>$#}-Zx#ver3=NFwz^D0UB~+z4jf{bGQ^-#MYzmg#?-vLXbF@NYTwohx-U+ z&;=-%l~bQ0tlDg|^q^eH$uWV!&~Bnp1HB;B*RQN+tmtbQ5QB3sf@3-nv3jEPP$gE` z7ExW1pC(w(dhaXdg5yGseO_08<@sdT4xHwvpLE3OQqMNy7Yw|Gr;cTLZG8z?My1K3 z%OsWx(lgHCz^8h0xU%5P4YWl1~?G6i)m1w}};#zR#4EU*Y z6#BGAngtixbODZ)Hq?reDjp~#&NpGt-H=N|(Q3X~uXg)M^^&d9)S#N{p}xH~92gz9 ztSHj!^$(yG-1lpev-adM2(6s$DUMHuD)y)z;1?>7gElwOGK2u_TEba%f`)-Ic%xSV z2q+BB7V${MqgU}}2Wns!2@y!$pE98aNkyX5+NY&G#5zBR2L7X`V=A|U&A*PR{Q12$ zm7HRA%_yGP}w>Sziu|>Ol`% zPZ?#$P|u^a(#>hsR9`5qFa|nH78(FXTnAGeW4{q*#ErkRtLH3lb@cG6^A-B0jYR0a zkS1?xQ4me&HHT@aL4x%%&b)qWz_dFi4k4TrfDM8+ulW-+${TUn+sVNWu!@PXgyX56 zn>o#6^n|w(1=E;BDfEOF!p)pkbbKz9{^Um1Zqd${taat;vy1OxnwZHXYuZ$0D53}H zZftiCf-59kKKTj|SJ){Cq1ai}dC*s$g{X!-yI&YwuNPJ??Bu^tZao8FS*}m~dfE3v zcTr;bfrhUU+FL0hDRW0+H1$CK!1l@X)Nyno=tHr(lGlfm_aX<_>K2QquK;vj=j|Oi zy5t6*JA-DWBp6bXx0FNS@u{+jc1MiRPCB*C6_fa+w2007?YJi1fUKN?&k=6AUfQ^-aqB=dIsMVwoy0zGn ziX?Y88He(adErxhC^f`+F`YC0cq!*S+igWUyjqYjT{-kpC zzz6(L>kv37Z`gKS-r=!hV#l-Cy=}|JCbCTBgS_o8+afN8Z-5Xs{(Bf9(^%a@C^5Yv>@9#wNUB_?fb+kh)H80l3}b?UqTfjAhS282rq{y7 z-FlV{qmNHuONo5lM-f=#10jA&KQ&YDr%Yyi`;+;_m9U?vhjSakj+DXOJZx#j$PyVd z6>m@iJESE`d(e`A+sSouVI(`DucHbJC0EnX@aCKX&%Dsdkvru@Qo7!dMP6A``|{=F zC!-Fd)bbT^l>XVpM#{?`zn^E4fpc1?PKJZ6Rex)1sp zfB&&CrcOdfmUi+iaY@!TuxCO9A<^P}hLjB!2DH0d5ks%TPlGW0B43GsC+K#?#<~!I zkC7k-N}P`rVqIl^oD+dzZTk!u?fUk2L=dM3ApUj2@{9+RJNsgez(oY~qkEWq`bk{s z3Mviw?}-WJGclApM5gnU3VU@k*8;yB;jJf3Q=2I* zthP@;>#9h z$G)N05Pn0Q1V!XI`3l(d{-eRn;Cj~F0E8@K#sQiYQR+Q9{0+L}3pXUHtq0UoSR2}@ zQ}Yo*9ienlg@kaa*FLu1qCD>cfzO;C*6Rw}CA;g*oiu23H;;@V>oKI>v$fa~#fPzx ziyPY(aoS4XJaLL~mIruwCYhJT;NAS+1YaCkYMFdz&BC@K)72g+dwPVtg>7e+WvrMN z;wyIp_D6+ct(U@I-9Rtxyjo>16`4@D5<;ET2C3bM3#{KQ@nztr_E5ua_`@#o<7=JX z*^7UTgL9-8E-%9`7WOwG4~V9e+QtAVhQY^~kV5izwl0e(n*Z}Rt|?d*l^iQ|6G@Mz z4W@||924bIm{jj@93#YlDzlcX4N8Gez}utrfEoi$IdhTwV$odXx>8uq_?;x+O4197 zE{}bPWYnJj5jj^!b0QfX_=kT(QTK@si;ldY{Fs1naP!O{{Q|{&1S{@+6(#k+^6{34 zM*-CoVL}sLsw|E^h2@2a9t_bf%0Oum{l3_4WTljBg-s^uErQ*BC(jp15l8M6O8ime z{zHGtSKjkYk3!8{EIef%v$a&t-QeX+A_U=dh+$+%QA08pw~SDUuz(-!?uV{s$8vR< z@e$1AY#YR?@PyhVGu&84Pbgzz!5pR)Zz}RfCZriDPPeJMWNf2@{)0vvE4P-EfV#&P z@ChCOJ3RR^;Up$SZ3OtAdaLV&fmsiz&UAF>R#~HiMdN|Wc7h~Sxfd?=HF%mV#TnIj zq(oyc(8`V}fH&me;@UZTWRnymT?@LDIke!+vD{X>#()CHz_>{xW9~a`(DS zyI7B|Y$x$!FM(LjWy~}v}{l>(4w;=Y~ zk^NV_`#irkfQT*ZlH$~+kA{6mub+v%mSU~H9;GoU?#7=SpB&8^{eKD)2Z;qTpOULL zZ0vqfJpVZ-@!;vl5>N3oqFrW*Cb-JfaQAH|B|yO&fb7=Pe5z;pVYx5=0} z)(0um7A?@rL0wy+BP*yD+#6V@qcLYpV`$oye?zQ(X^JLJ)1T>gyMESq#s6fVxMUl@ za0P$sjYVq&#@o?LT~T7KeiSM~d$y+50P2TDUyktV(HYiSiHX`Lsuyh6S~GZhb-(K_ zjXWEaNq*|4PYJ)=BA529`k81rWkp>hj?=Pg&03S=ny$p>u9v)VZq+ev%lT`TC_=>Q zYoD)oS*>V{eJj{jviQ!4#(K}&m@T)?^?AmsUqJdZ=6pOlWyTxS_bd9ZFU;h7N$~jG z`Q$~*+z&;a5zo2xqjSoXjn&3!`P%r{Jd{Qz)HF?*Vbl5D4|xr&6~QZ4zWs9KBMmte zq^*;^er-*)X5SY|W6;?5G$v1t7ktAQB;;-3jyC>=f+tZLk{dN*KeTX9?65D=%Lm+T0b3Q5r8_)0H?Hx@2Jm*3_4$sSoR2HvH(Bb$x5S{^k{w zwK_(LxBr9qy~9njbEIkNt-lSv)6w`9==?YBw11oTrpes)K;s^xIc!YhJ45qYp0Tle zbKolXH_;n;ELV5zc*CXBTN-R_40f#3`3|ypQv(OEU?o90maJ9+lsxv4Mt>7Z6YOqE zDg3Ca*!+f`{O;!iTctLe=e#^NA>8bVp8uqNQIE5=`?uJ&iD|?eX0{VN?KyK6e3|iH ze`*Myd$PE(-d|-GV)#sZV#8YN5Ql56vtGF01#uXwG=RuLiQo5MH94WKeZS3&X=%7< z&G5r1hAuZutyj3efM#vqV!2ss{-GHSdoHuqwqdJ!k%O+ONn)~y#)2r(@ zPw!j%duM!Q7}Mam%k~WqW26j%$8RX_G?iPQ6-BTHp7U!sZ@V|&#$3xh9q)O=;LW#} z<~m0&=tP=~h^Bo?^I{rAon3H?vHg96Z$04XmxP zS4EIFS-am=dp~uH>IrB-6H_=SxS?)Z2i3bE3~!EHghu%yy|I!oiAc8>$?mT95&GPC zfK2e!E8E+U(+Fc<+n9s1uE*IHVfbd=UW8sx|KbozlK1_F%9EE(n&%Tn%%L3lDbis|6`C8KXv!9Jo}KOVbm^l++MPNr_dfBP9YOrHyl-A?)3r&n(Nb z+wrW^Db;;6u%ud4Odq+4sbvd|z}x#m^vIo*jYQW#W@SmYVILy#WIRD zCZQKJt$6~|5chy9M?IiYY5fC+K?WFL&P~?Xlh6I?R);oBRlBn=h3nnv$2NX`q8J)mOf~!@}L~0&V4PPM1Pmzne)XfM3dL_TrqTDBOAo4{f z?jiXD;q(GZP;GhUscX3pfW-o$-7h%dM`kQ5TdiLdoP$=Hc*!y;8B&+D<(#=xeGV*2 z;VhvBA_6hnoOq9-_6QS`_~bM?#!(Q*H7sLa{VbAh_puLujCgvsYyyUvX@UkJ?Y30U zb_=}liT&gKGZv!6R1?scQo^07RJ6hwsV4xvTLmLhAOyWY&`Z{(+QZl+j?E_6jg;Y! z$O*x>$Guc`n~|U*%0s`Ujuy-3!=*HC0yVU*+!*kKB#4$-G zQpo!=LyC=`q-DrRQ>Zb>H7nIe&@+x~%)->NN|?+daDJQwtA>f(yAj8fc{XmXtwJzi}H>2no3sP06)A+DJ+Dh#%*qpfbOB7+S> zXulP?b66S)J4WC2I~S7eSlm?P?k4p($IICerVX}=Ak=9QhNDMadqYf8&bAiuEqt#k z&&cRqK;sntl)4t-N&>>93(ImV{1_uk03b}iDk9u|UnTTDqYxqRvn}AwMfAg<2g2YL z5={MwPN|Jn7L?4LY0S+2)=F>pi8Kaz$ ze}Zf7ZV6|~8@D2%z4{%IMW@sXUnsIX6__`$#^F!rwVvwZblig%Rdu=+iFmV7yEcKJ zK)`#%B>2#^mk~~+?x*-~KqA?c{CyBTw$l%0#*dB8p?vj$h6W5NPw7)B7U3Hyg?tO# z)jW7wB06DwC}7!%$(Z^at7lc2QlZ+OXnb$V-;vh*oj?1X49C(0Pz0%)L3ht4_U2;3 zXHx5$Z#_FMj@7&MCGl6eWSM_CD#%z|fB0?VhBvBTSFF%_i@tVg^O-VOfBy|BIg$6) z0*m}KclG1D^jbn*iulJf+2ir2!j=63;k@d!oy0rVS`IT%Ss(sS6xqPPKvZ~i>}ADL z)QhXXhX5A;u1KLx{6E=GTjho1Z%8aixb_6&znxRPJw6<$)!Dtm`hT(a?%`1Hd;kB4 zLTH@jG*hWm4pSj=m?0sFNhczvQaPnb$i$ct<$S2DQmjlU6q1|@&0rjAS>;S9%uMB+ znJ_+?!}`AMd*6Glwbypvd++ak?eF*Z{p~+oS69~=4^b4MTA zoVs@+DXK0hN`jUB)-~nJ2gAk>dqi%?-Shf%WruWX=FJGz;6tco~SN&&xazO3<^o> z0CNN6n>a^c@6Vi)wzOWq!lctb)O)0#Fpbh`E2`X5j4sSPnp;B;vxl<6iNA$_spy>_yX zE$_Ht^23j5b8&Jqnh)8LbkKM?xt=IW>l4Q>&$?l1?TKQbU!UksgPcSiLmK_1BmsPX57@yb` zyRFdY31a0NSA=?P6umku-rg?y=m9z(nX`@UyPXMaON|rMJu-W4(deDXS zaq4kaOX`(oiS?H2qWZ-UIdRLbOBiceCWzSc9nQRa;xp@aZ$@<4S~3q^!irYzdn3Bh z@WbHu4obQ8RgDM0F>!^Pyog!pEs?sJ=rNz=ZtQn6d8B;dl(&T3zPc-vop&QKuSl5> zkFoOC1$Aj7hsNzeTU6`XRM%?#_zCuUDA@Z= z8u~=8YOJ;%Tx;P)U7h$wgsq4w-RotQ%noDO$TMQ+T)!g@SZjUxbVlyx2f6Np|Ei!X zdTJTwz_PvOBF{by0znc}c5h&mbZL~tL5X#hR^kOYW| z+Hv>NrS@-_@FZ#2C96*-<$jcZ-DLOm46^pxr@P^fs-p4sZ_L;9Yi*Ri(G`RrwIy|| z5xW#ce9xBj*(dh3d7ro_6`q~5+>IR+8}b5qJxQ_&2AW^a=(3C218OlCS@Q50c#(aQ zqMjKjZTVr3M9OxnJKL>Ri-;*NN#$T;%N{k-%ihbbWta1+FT?MW1UVW@&a9;)^6fVy zKPKH+BThtG|L(Z&K1}-E#cdH09>C51gyAvB4LMO;lD4O5oVlC3amv&FpuKsBcG{u8 zBc~84!AvCn0;@@lKLWo=Kpe2sG(D3xzj2`?05XY7N30uo@)Eov_lTN=PkF<$HPJD5+_b4BDvfb= zfu4*;t#pT?=>zkzDQSRj;3@nWtx`dK)_oRx*Bq!XewD*{*6_TbpgO?Lu4j7=v^E?I z!(A`IoSAv=-$dHdUaN#i9T_YRJ;=dKy<|r>X6^1-M@!t(6FiAgf@+R-heWs^=GX$m zsx%Km?BSPv6^WXDU;GQ6Mq=zQVEj>67m($t(_1-%4BYf_MrssAmh6k`FC3v>XIC*l zwp>S}wc{}7eu1eI5R=>s*aYMn{uSSBqIz%W1Zg*WA+k1zAy987OFe|5vyEc1bd5CK zZ7VoLPSZG9=1(2W@4b(Jjzw&IA5A$J6)lu!W2ZZK=ZOS1?OFvbmZog<){f~CN%Nuf z*H}R=3=3StE}<+G!8DUu^=*NPy#9%!T*|;He;}Mjt{y{ga4mZDBGW#hne`bjPio3i z7CfZ+r6yBWkQ6^cPmqZ_gsqt8c1@W6e7bO54^_?gm>O?_EC#xjX|eXA}J zsYisg6YXdf&{z=0O=FU!AzXu6%*!CqTXL8KPg+8)p6D@>G%s_64g<9NVjEzU*5Otg z!Q(Q%rn4Ow#&KW3ge!RT0kme?K2B`3xeCt?bg5h^(sbEV16f=l`(I(n;9Wt^;ktPc zd$u`Cw9gPJGf;9Jb-GWtLU1Yk20RaKN_tNd=Yw}Rv3br+{e2`iwY2YWBSnUM61Rq^ zA-I#x#x}O1BT&;x&9akIc|%{bXi-y^H~?Eyj?XI)G^|{ODef zE8kqq&~tCjQmKKk9Go-cjaP4FpjHr8_kNfde}pacIDTTh>X2(|kY<7>mh1BJB}}kV za32r>k=JQ!YB#Zpku217^h2V}%QPmX1!y_|*f)R{McOej%AI9vr;yLgNPvE2cSZ!N z$&}VGZ7Sr{YH~Kvn|3KQac-K6U+5&H;1b?g3>BYIdX*2p%U!#nBntC<;VMy&vz4k( z46J}O!RW+ua8s{%s4-7T==$TyVHv7EBvoGx;0}kG_clg1+#t9dM=JBRnE?uQ+xsE2 zCbfkm;-#z|=?!7jSbwjPRt#n~jjJz+Lrap{DH42C5php18ugyI7it6Prv%8E-O`9y z-OblJ#m;A3&Sk4}gW9RasCbe=Q!&A=QLVQkAEG#O^s??NC+7GMNo2J)EQ%)(?F3~P zdz}Y*Q(Lg_=_BxH0EJApX42?bdA=V;m8wis5?rNj=kKR9Q?=NnT(Av}A}dBCD`YM# z&`B#g%CVwjp4#oWy6PfNjharQ2(PArsZcj`c7actZ-?ht@8Z`dBA^3!<%J%qDr5jU zals!zUggS<4X?(W8zTl5Xv#B(ZBj<}L;0N9=1`gnY2BEMike~Hj#do)KHLRl9ruL$ zXU6Lrjtk{my|riL?Z}6LZZSgGOA({4CcW!JNI}t@n6A2kqa}^01J;eh2;s`Mkv4UI zV{%0c9T$fybBbjTGf%pU-p+ zfy+S&xV3EsR$Pk%qO%IDlIbfYL`%UnvR6s}P#Y`1F;#W~3BMrR^`7qz;mJVeqfNg| z)r@UIh-HrGXFeXBC|0>vYoABJS#}z^66Ny7yC!8pMqpKGWmS7}qw;KcC|_QMTLU@r zv;e1PIF7W1X!`=XNJZfz2OO`@m&6irOvwb(!5L{oK;b?f}flkL1!DUF0*VVt(XuyQEworQ|#XZVYO@wO%8|yC%5b zjb_AzJww{;fY_w2x?{=g%|O0m!bgMkmRj!6CJoKHLa#Um`#7Y+hfv!7ZsK~zE4P7G z+iD;x?GxNbuja33&(x9QX)q&Xd?(P0wc$Bs4vt#PFQp71oa$AZsYfM|HbB1NMf&*+ z;D2arA8X6GoaXyrk{hO&;Ru?5WARUoJhrhs`%cC59lRmRm0I#u%^g`qdB{KXs#Cau zuK+Ci8ZipOcS3!r6p&c7>EsGif(y$oW{&bcmZ8!>OYmIVTTEIBc^&Kq3=JSX?9bzS z)b&>}+u-e^LDLmL*KyY1KuG7<&`VRrctI8G>DcTSQ~t zE2F>YGFr~A{3Y>>=JHItdECJt1!rMYUFt0 z(V!e-z#k@VtL>O;yACiHu`Qb6+>`9K7__1w@GDm)wV3hXz{p zUiCcsd8c(CUiZdZb=#@@c~BHGo)Y10nIOxNw(`$Cs}ke@=~M|Y9ONsanzMkjGoqM1 zxt}vcW8(TzfbCBc+ent=TdB|$7&GF06V9f-W=%8amleI{;w-Zn5nqG|r}Bq^^f)Be zWB+^6@9Z2}P5??MX~)y9XMtwXnK+^^&xZPvMHz@0Dakq$jc@9JmhC4o+#-FFW6577*WXbhv8;DO5Gm zTjVOD3llSdY6|n7hI3RH+Ec@v29R;Tyi2Vuu?T>u_?I}9_Z|^)==7l#omBXEWv>jo zuR(3Tn2*>YEi zNa6rI&Oh}SWVb5{_@~)+P16d@7zzWYXrFz;FXJ!)%fL4{Vj04hk2c;pM%o1hGTY#O zVfok1tqn7X2<^TAmhsWyrIfz;QFc(BDQ@}-$G8Q#8feLogjoh^ssp%OubkitH$jJ$ z?=#-xNyfxg(3NmZY^^cpFgF4k%pM<>ZS*v)HuX*|DK<)!(V?diVl5Nv4^3oDrzh~y zBHzPpNIw#FAPe>^6EkJcuAWwJTGX#%VxcXe4YVk-Y~B^B>ckzjiu7cnB+G0NCkq)L z#H`jOZ6cTkAT4)QA8ol_fsF@!=f}#5Z33AuYcc*|ixGKg^C(bT9Qkqh`RBf-eMLgm zx14vA4424SUmk%V&I%ZAf+sEEY@|@JeO&#Q&{2O((>{{cI9a6jFg&(alxYlAJ&2&u zdx5qKtws3h-IvX!U8lQO(mnelNpjA(8m$q`^`(C6MBolG(~qlPI0Q{H?p) zA7+^W(fT}D)y=MbEEFetL)i&ar`@C2$}=T}tlw0ON9`K;K|lb;7- zqF*<-Sl?NL(NB*=7Ex!bDXQaBaR1m*1D*9{!n>r1bO#}%hMQH*>GR2*;A;Y#WT&NqYksur)1o$4iJh;q8yp2gR@+(#3sbLZeXu5%M_o| z;;7UeZo@MkEON>OK@w(=va2<$WUqkM}g^by`m(p6e?_?GcQyX zKr_d&OBH70WIcVu%v1@2wQX8GW(~!P{mB&I8T5#F#v~<)8|Z85l>Io8wA~q49LbSS z0aI>9`c&E@!h(k&&t??Fw3+JwTVBJS!FKBN4~LK7G<>tvp&lkaHiH~vI3q@9%exD9 zdbl(U!Ta>*af*aYmhDyY$Z(vb$~AECe;szPuiZyWo%@f1+fvc(zxi?LR{AE0u(CvI z68G%I11YF(LgW2)@<-$?*}}V)pMu3R=sr2Kxq@!;L>?wEEH+j=|NM_(XBIOC|0%&K zOsvD^8R;u#{#ojaSJE9(roroOMA#{z(p320QiTwXUGR6Bs$!i|c9mT#_i6loZ_|-9 zUZ2J-gRX>94d)AYf$VNS?7PaicJ;fK6jKJIG(xcq4o+_j4Db%oe`#^*gwDfyufk zk;lK$*=<>2vLTKxxue-WVW(5mr`uv_Er+fc+9|q^3O*3`KOiR7eLHP3-nEhPSo@Xi z-`4%aEhn>s>9lE_{?;hbzKVCu@@RddWm3|~k#7c_MUIMAN8(uGZ?9cUZ2vhj<$nFT z#V9-LqyM;I6z1@$2?%*|e}Pd|_}jMp@lK3>r>KuQU}(w54L zNkExm!1zA#T2Lt&iczK;2YW zkp!(}U8dEx)PhH|5ai{MNR3Iq^J9u1Og*NQWZj3;rDag5NCd zs-vSRJFc7-^AU|GT(@lmNH%S^U%a9^er(yBEjIT<>Dr#5N96CmbBGWp7o&eXf_w93 z{p$x(1~#%sff{f2l*WcRgLPW@7x$qb+-&duP3T_EUHiAChl64lZRBF)y^_)t&)iRZ zvrSSv<>5DB=!3QY#X=LwNqm>O0QUHkyMm?( zga+dN^T6@637@|i5TciB16fv=3rl~KD7*r+Bt+n5bnuV2@q5|o+8^k0n-AP&nLBRR zC>Rrs6n_8{X_r1IqT!6EEI4zvh<@@EPm^oG9s-B*(Es3``psdQ@LL`27FgrIz^22Z zY-A)bkSQ~j%)4)(elQOm&L7>04)@C2J(*?hw5%!t{2EIcAilkV8%@23wgk-%&%K}N zq6YRQ%|`x}29s1u+z*NJ0-zwqa2+NFw}$3RgSHGeCR;ZRpr&X~i0bUTk`an3-%xNP zpw@xYS1pt35Y0lc&i^80Z z4*|$9IB54m)giXCXQDz<`-2*8KBVsMs+D{MCE9!FZPbYT&}4WsTSJ%G4dJz=-e9yP^b#|l@6g1Unz1y1v(s1Z*> zdLTu|5rUTeR8!?N9@d9aR&pJr-QkU$lW5I$7RPXC9aPL-psVE-K5ZMx{3zwrGyT-a zq~vpHNAPRZ!w468JC1jIfv+^{N)u#L_4xZcNGgP?k%b0S{O~Gs6}H13U?Qf?6lRWK zhb~*SS7r(1n{Wcw%nA zM80?6S_;x@jadN43kW+|nw_z6LQ~Md*&O=KZDiN;KiFpE`=R>CWSYx98MEiWh(t^!=WY!0S@vKqYD zW0~o6T(TOskeQ5XsqPkPby3v;f;T;`gs;xfG+-=5oBJ1d68hY`ic2=OXZKFe!KHVX zR(rTKsU6PTTLYWJU!{nX@yXr{?Ut?NbJI9xUOIcH9=2Ji0N@cHMOn77>@RfY%He3D zYB;lPEfgeFDx%35`qL7r8#<`!{U0WTo502?$7%Q}Q&sQm^G*oZ`CWOkgS1Qd6S5-Z z>Im}MKz?+sHz$Z0+c1+#TSY$Is%dNC-we+kM{kB=ISLKKzK{{moZv!tYfA*tPWh=& z*>zL7ryP4BX71Gmax6*hW!JoWF|r|S)5rTiFAl>n0a7!ORg+)Sv z-r`k`apO`%LQc5@`MEdYsRx;Hw66&%0hKg-n;2e$-U@{_A`qP6YkKCEAa^_yg{nQv zj&4S;XH#k3c~Isa5A7E8%2&Nly*uXvOgje~GmI*Ja@#ta%d_GWO^5j#^I7C3ki3KT zgUTmudX|G~_6E7{YpfPGrTpu_+8Z^En=Yv$H@L8n((hfN&?e6fv!M5rZ=HSI%~$Q! z=PwkzOf@NCnDic`0&vrMJO`pG@FVlh6Lt=56J!cC#>whL8|Dlsa;ai`0GE__p(n%< z*XSwQ9fFh~Vj%ZIf=V4ye12GTJe`VZuT{-?n)`TSA>A}-tM_9KKHgW(vcTfJdCC4L$~`G4rShqiKxyL<0ch4sF|hbtoHIe-g~XW@~6A{{y+qrbq7V%w#n5d^yt)V1J+c8pIP zjgj@~r64P_etKjWX~kZHyb~(W7_X z+@*wgj>^Hr&*Oo#&jAk-kRo?pwkmW850CB7pRN&X5znWrhWg>a7G2>UVpzrchZloa{Q^tE`C6<537!}qH|++1xjLakyIM!&v>gg+koDfZrcmpf*;GG!#h^v%W1_#{F%cxP*jQvgoqLQunrv>sDv3MfH%o|S zF2vWO@(8^B-Fzc3LgJ~L08Cyy#Yi2l473a7PNfV7d@M!! zysj>QGi2N;14I|`6h4MQ79(tPa1^zD$vKwIM9Gbt?s^pXh(MCD4KExFKqyU90Te{I z+Ld^mXE{NZAQ=!A=mXQS3ALD9VmL>b$vX|H)SJSS_#rIvbS(SPqrRp9EK5Ir$-HqV z$L{?+xbGnq4mQbX+S#$Y+UuvP%>3_FTzlwy0MCec`WcjOMs!K=Fth0?Z`CP@rpVPA zum_ti_o)?B#vBeMv#>E3m0B$(c8$mTOU;mw*uQ1n^Afy!pC!EpZUs&Je+?hc~z}< zE>uV#GU*J_aiAzRbBQl#J9A4Q$~!O!YAFXW=kW0eu%>wX_-8=ZK}k}z9yObRmmnmJatk`m$sF(c zflK}F9-h7isjG-bINr&%QaO{0mvdzV)&3X9GF+0TO1!&uGu@*69}jjA@jwS))WLV= zN&6#9E7RCm9Gk&`Hw;Ho)}TUaunCKx9#tM-H$qMdMydE=WO+&Iuqvo>E+PsaO=f7*KVq2pq^Vlf+=)d*6`mQP;M`feJD{lL z;rN0ZCcbFe+~Lt1Uw_GG9HrvQm%g9~#!u7ZnDq znQif0S6X%_LZ{&lb#w0nzylub{o1~^$g9_AFf&32yYnVDvH#fGZ~b-R_8fN>xfz!- z@Hsu6WJEj$#dG`{)y`#Ppzjs&?KtUi=H8Ees}BS@TmpVHp-dU;gBi_@AF+W~HWe4z zh*eb?8$Zwec%sel$788XvOILTuW^?rlC$Pz_HjauHF_=QY+h<1Zu|~ao88efmnw^d z6JmybVoFpLOHU3V+e6A8PS^*X42e+K*PToo3?AH_P0xw6n6(xZV*F6j1hrNMIfmwo zO3gw*zIlzOx~Q6CBFp#f$HpTzr5>*8Jv8M{ zlg(`}-&Y&*4?RWK!KM|?Y+7?$ zLY87z?O{eMJPvJ+UIoQI@m-`e;bPmNh2{afXsWItitmR!10_EqkQ>#k0OMk3d4oe? z|GWGPeFi7q$IUt)E@SoLvZ1rN;Kd3MmO6Q?A8jsk}~>2yy~wo1b=3}CIr8!_xy zz6BtIx6yYY6hXTUe>+u|Bc4!eC%Er*>SSM4s9IZGMTFg<-g!;;V>Pu`WDe`No$Udlk|SXDYTq9I=V+_))vwh~)XBmP+WLf)gfC7o?!*&%_ybs7rHp-MSc6`=- z=SwTLEDD|hkWHT5z#9RR8rBzAfmfNts0&ekL%Mx@d$6hesav{*n+Wd!NZze4QK)?U zgfC5+c(l;p#jU&gPA(H0s}|DJ6B{|-+IkOA*^Zz1;0A?FIz{IHDa>^YFnwBLVqDxBh$p{zUJ*|I3nj4=>jq{FfKh?s-S2MP~R%>w~Yd!fOU! zetG#a!~S9D`RXmA?bAYizH1rB_?kKJ+rTh2OOmKapG8z`?Mmh#Sq-rJ6bM?D?2nn_ zLz?aU^(8ispB6HTkyc%wZnfBS8T!}g=;^<+s|2EQj~kQ3EzmrNK4Kv(c|w$j=o~YK za(2mp&8eqDL?cdLe62!Y@$tb#%5(RLPD0S(u9TcRsL&wveJ;&pyzf%Q=lH5>sqNqV zwulal2p)oA3&+Qr&eTI_Q~t?BhFyaG*BQPYyC1!gP)ATdQQIA&8j|e9I>N5vws7?( zCkATY7WxSzGLzkp>@(8bdt%7l{jExZvkgq~@BC4kUrD6l4jvS5!Tht(5`a5-U(z_N zz%(X*7t?rb^$TpL@LtQiiJ;-U9mD3w*59lhYfIZqd}(X2#>ov~Hxqa@a?L*Tq!k8{ zEB>yQ?0Jpl8Y=96b%{H;QtzKYmt1wnURJ(%RWV8Y-t5*mQC{D8wao?7OX3zgcNgw> zJ(wr>Id z^7?P(=j|f@L?gD9MshX&!|R}R8%5#BmvJY*95lF^wR5}O?$^s(`~HtksMy8$+Plv< zS3kEzo<40kuw7CC=4AS#^x~QyY?q@1h)5@`kCQmF&Tj4N?IZl$oA<2hCw7X=!#b9| z7r;A_7d&r4sp9AEKeVlbf0X`o6RlkTp3i^5d@fHGuv(VP zvV3>NvYjHobE}+pz{+QuT&_Juz7p9{Cdcc0x9_&qsNm++77M?byMJlK`!CwJU}gW_ zq02-Y*#D6kT^2U< zOPo8u!0ZS1o`(tA-`Qn2oC&7u!?i88{d{THI4x-X&yhEz%}HO@N8hu2w~Fw(USo9K z^4{BuBIo`G57B>3%>939MkAMh7LZ%^0b+UUZP@Q9QCnt_LW0BhdXTX%cp%yqEey8- zToVPppUJsv*0P{a6*Q3Dv|56iIDyW_LL)`w#CiZs1BJ^mDNDKK;Q8FP=v!JJ2+>}- zOicqb*Ihdopm&nZ7X#;vm>dD9{!TstzaQs6AcT`!ngQeIMDVa+&0?fQrO&#T@!2n8 zbmDLqpmWKa;h-D$bGs)Kr0q2+piRCy88j~9nO(tpYC6Mm zLnr)01IH069lGX5jN)x(_njObr#XTbXv1O@eNidlMMbBf<1YICE@suUvZ^RkrBm|n z)E(Y}_oe`|yNOr|Bac=}$tzq#d!YbTF5uKS+k&C-Bv#mJazh^{!235 zEJRVHMi86x3#=XK0FpKl`5^h6%Y%dYd4Ns_xlj@Yo^( zhnTf;r|kzBsn_<(ztMI68BZi^xaE3GJHzk>ROR(X9xJ*(!NIm{?>pE{DWo<;rp$YP zzBuq#x7_j9f4B(w*W+&Yhe7$vM5qt{Zq;C(&EXbB$H@??97xQKqJb*-&Y&>9Dw@nb z#VQf{15L*tMukXW2+2K^+w}0iS`v0#1SnWqMxa1aj!OY5bu#19Ah}}~8!wbL5BUZ5 zHUjy_HIv52v@Mo>dx0mP8l%mL10X#c=L<0Q{yClVj#%sVS}Y?##hKa(vd2D#nRrzw z)S>1@*Bna4-N*Dtx1sAG?X{tEU$B5AqjEe-6K9`zw5lf6shg1?e0!=oc@u*Cf6p!X z{MAhUe{1pnVRzfT5h6zg$G@L2ru}Xe8UFQ5ivM9t{$WnOW9a-}V1*c+n1#PuM5({q zvmm?thdqh;!z%mFtuypj1OES$t@H17_Mf5ue{`K)f#|Y(f8ZOiBa{1kHX09^NldS` z-jGXC2{>o9K?QazyZxDZ4E8_c88}LQAJG0VE9%74*Y_9@4cTEu32r<4gk~b$2J6pg zH@!c*Cgs{w%d=qTi=UM;h4NFYE~XrjmlC@qvic%C>Edl~q>MSj`MC3sIB3SN zCU#5FFgyLELEMKYX?G+(txz2Kb+?ZGOLr>@I&pF38j>2pL`V7HFrL%f`IF*vaArwf zmOMD-DtbfJzL}A^7}^N_XQFJ=(<(*_yYybd)3R(FZlXFOw^Yzk+dh9>ze4r=x62r5Zo&1^lOrUlYj8gBO2RNrmj@>a+H-%)NU{ zm$i0!#Ss_o_l_4<+X-_AJWf|F0wK%f55tQZiXs|idya_cp80+wYvbX%jGb7|$QAZe zBH9UQzdkbiWPU{nLcl=SwH#_+b%CPm;Jv#1W+B{Q{e3Id+pk}5l)Ks1S+v|EbD|^Y zh}=S+&KD-%Dr#8(l8mQ?;HhW*1Bsv80n%lc`4VIy@lSXnlRq=391J&E$~R;0NiuEB zC382 zMmt`Sy*VkZ?*Z&qILw~;@V9PVzq!!vauGMxY}rNtQ~{)*GY6O86$*c%9{E@wI8eF- zh_P8hSfyO~GZt~5nu_Q_9&_AMeHG!9k882SJUd+4tEx{{_Sxq zc#2_h>B1q}+@YR(ROTXm{%`nNyc7Mx%O6gqAuU2p?Ouk$*Y;0W% z&@X2^m*9<~f97miE$5{7@Hg&p*1x*vS;{SQm5+Zl(~JMhMTY3Xj>Y)w-65)ba*5xf zxI|YH`?LfvS!N3D#TVnqlRuCbS9BnUDZ=U!VO&h4)^cuf9m`E2ejZ3v`nJ;MOt$>G z_$|dL;(TtRP^b;{=%HQjyi@YkI)*rBZ5m~ubMwYNEiK9;-=mtwLi22Z)1WNw+7G+4 zJ6mSQ&gLaLnaUC`XZL@if0ncAdi-1WxZz0_*er+1V;2vKU3x2dP^>iV%O1i_>Lt+D7yJVEtI?Zlvn&x;G=o|DW2Uzb^ESLmifKe+Rbae;b7hKM$E))i<|_8r?ol39}=mlDoop zna(D8)2uwTm(B1Z`q>vvfM%I%D|mM=x!^F;LKt)pX;IpxG2aG5Or98Svb}{7gSMI6 zwy9`~>OBotEGnlIXuJdzV`HfQj#oLQ3SLkNldibKebNb!Tl!@u^$xfq?;m(rp1x#)o~Y5)VeYA9#lfe!J8EB>a6F zq=lUd2aS5Af(J%2YZd^wbvTfPOD!SL`iIMmQ=EYs&eXsB$3&C4x2G*UeT4=mP8dpE z!9VgD+SH!MGdSS}#nz?YIf>Q7LMAnJBuTdRbRlmy0bVyG&mN89D?KFrM1^(EQC1Ti zY@e}5W9MF3Pz{NGw&h*H_tm6ybk^<*-4t*tar9i8*&9E;LHbBvl33_x;A5vMJnWjvCbSN88Ce7!dC?W{rTTX8T=lL> zlFh?i?7bq~r#0+Ju4`wlGqcqMG{LnFvq$4E7>#Rl&UV=` zH`b;$4k!@J4wkd0j-^*FZt1G%!O8hj*LMfMp4Ukrtrpxy%Vm0m*CV#F`(o3M^?nf) z;Od69)0X_|uuRbDPSmL7T#n72tTw*J*J5O1t_hXe=OTg{4)y!TGM#D)a8MRwW=%o{ zc_27k60M%u0D7&TZmxa#f@3qcpcb3Bo-@s#C7xt=m*EH#?>|=1l2Mz;s%0Ob=N~3i zz98j(ZY$c0zwmQt!i96g-Q$a~(_Rn9-LkGf0zMjdGIJ(f7V{n#`5r?e=2{(OTn_ai zt{Bx!+x8B(+@yzx6gz|j>X%wH*l&KU+om^*UH{P%QNb#SBFmIx`&H<}2wlJo?#(jg zkdjAGYj6ZV5AxwmEJRkA9NJliPx;obd70HH?~7d9Ro_15snk_#b1HzLhjC47axpcp zzNP>jlESSP{76)1pyb*(K}`r0J7Tyo9CsZjQL%(FpierwxOU`n^mCC3wd)x;%yBB4 z=k6cRph>-lau%AxK82z4ANk1-j&(esCcGZrp6=8|^Lb_9;&BP7klD(8)@~SH(F;VC z$i}eCB@ZAsmSZ)L^wwJE0)wc7RA z`iLGrS7Br9&z0xsFnAu>2Z)E|II;A}R5@3^dAGSmF@)l75!|mmHG-AY&u41d&em5B zX+exaSHjIxDa`Q_u#?*?`}Ehvj5Y@Px$JWDXsQr=DlD+=qbEUSY;Ojt0SA~X zQzN`aNdm@uCMF9}*~|A}C2xdKjYxUSIZaf{(6-*Zu93cC7w29*{zcTxbZ~LtyW3vj z!Ft-2rkiebZ2tbyTks^UNs1K=QJC|b9N)t}TbZAs`A4Y=Ph$hG=;xx7s@a@T?vwGu zehO`5{U7z4a);jPb?x!K`Oy4p4FjJ_-5T|+uW?A5%2BG!EdcH`+4!gYlkDD-mE6MY zv0QrC=KL#bvRP(vUn^vg^anns9#azwrlf>b>Ao9xVO~6dcf3W8c-NyO2z?B6=~*_b zE|8F3hp9x?lcT{>_?pf_CJ9ZtG@$2$t`$9c=AlgZ+A-2gShYL>ug>f7@^=cNrxwp@zDzQmo$zRx0`bUvy!-*Iis;DHts80 zfxiQCXL_jDQPrbedd#gOa)zPk!zv|D{Y)LJp{QBx9d7Ty*>F91P0VM&-b2a<;3BNmCHW5a7i4O-`&8+j?Hm*ZDq3H0)-eH-iBnvL)}GwNSi{!Qz}EnhF+hwYqIi98$VIqyU_-%%9km@hT>@2phpV{;9h>Q^p1Rth=auusP$WPGR%AuR{`}6;bm={<#>iA{AdSB0`#h2>tr_I;e>5J4Kg)p^jA8~LJJNuX&#g;5Yik_l@M4n zPq=9T z5%Rs}J9v8lAIvlRF%<=`AhMo+zW1?>lN%~E|*#hCu8KF)7 zm}ui66HZS>YAoLs<8xn~p_EeLTdcXKuc-?{Hk<0)zuf286XM};&|gPz$O7eY{at)f zKXJAQGY!}(vOWmiKEV9QKhD8kSL$jIZtNDC`Ya4A#46iJp}5K=L9E9{GMn;YszVLT zRrjCI8kBBi`B#fiOEDv?8%;~%%{?CbFH#cE>lv0P`Kh%vV`L^=vNm>>^9E(?u$}Xh znQdy0`o($g{nU=7M|hngbq%LVg*p-=_>m{F`)ugmBjew=wyF0xKo3lFJhOg&Ow(9|Ffwt8Xs3L+F6CTGmE3|(w zo~0Cl-3u{te8Zw4Z=14cYzNW2?+)4|YGf?KD7m6I@8+X>XnLOez@Uwnm?U1pVc_iO zw0q76=+^8$U~nn3bGBJe0jm5;XT3}J4PpAOuTs9)TJph)Ri1e)LwyI=B#(CGbIQ-J za!jg1lOxfd)m688C%*Lh45$|8XAHS|+#cGx{$?)WW4qYplha=vyHA)7Zbi_5aNgpgn8_r!jvI6w8wp_6&oPrr%594j+sh zRn|@TMy)CtIO*;IM3__&(mUXnQ2$Dtxp0lI!D|zDe~N5)ukv z;s-W#d9o4lh{W@iVY+TE*DA1aDsD2@O!aaEnRW>S4JLDim)P(1YA@|85)rH*e1HJY zsF609%Ro-Y=C~avwPM6D7f|oZ>W9l>V{2V<-xF|kNEs3WbQYF>O&qbnSVESM@p-;s z_!#0*9Rs?)=HMfdQ z2R5*579(}a>%p6m__N)H_py!soEAq;OFbnz@oX-=(1WlRm*?#*G8Xt#=aFG^Y6%e3U@YHI2Jm^P}n9{{L zJi|&u_~V@Euc_IIn)-H~Q?#=4)5&>v$z{_&Up~3VR1?pHPt%wl05Vm@M5WS{LZ-bL z=3o{fkHSezc@;P7G5gEdLpk*lQQQ#UCTpf4-B#V*At*5Wt=EpK^71OwC%1ZylgM&h z>GWvh{S8=XGZsATN%^L~<;n8#g8RS}A(UIvrO!9>^#;-E_ONuA9I*6go&o8n0tWknX0%D)$>+}C+<6NnYx(dxshF2 z?8|tyrh&CD+U?fBBuCV)kLC|r@{mn5>2i1y!lGTML0FtU)J+T-(bVEbKkuqlW$6gc zNlaH{wkbmA+FICS1A=#@cxhR`S-(n`Lw*+r2V00;k80QJ^iSX_^6s;#ypa?N@Q|#; ztP0VIsqQAp4-g}|im7V26R*Ba+;R*Z@gas~slK@7;k=>yi*InVf#3G}2z_3``r9z{ zp`2HT{}^Rc|3-j?34siKBpO*p1Ewh;=TmPaG$XihZ$m*NI0QkB&V2-4L90jxQBwBA zt+LP#_JLvlFzv)!{xbgcp%(|5nGzMW@Z)#4NgmjX5q6*kZoFlp8!$&%Gtuta-!NSO zv5s^{*O3A%s|1g9sVW{s`BKj9>C^@@(7ovaO&*4Vyg0aeuBTiNA9p3+C*uHhU2nNN zPT9u0Xkv%&kVZ#w(b*oY8a2P;ce?KMI%KxIU>eRu5B#9eVw@O~tYF*4+LRx?zc-ko z;2&z`is~=Y3)X4nN$79-^nHMuo$pq^`pwoWUsqNnzZVaGX!x8U%9yI!R*2!+R*A{n zY75G9hpIds9PZg4blBA8XBvF6JTsttlF)No&TDU>%VGM~D0K(dDMOyRtFB?Igwxp6 zRKO2KJ5&Z@u6**^)H6Sr(@HKl;em6KUf7v-?BBkW|MsPvGy5j4y-oumDR%YWzLZnITM!>38m-;9|DVcO5eI%dE7tIF#EnOYmL&jvqDY-ZmJYOy?=D zt}1`kv*32{R8O@FoPqTW?(q0?)Mds&WnRqVQs-=rl`+HjP;~__&qxQwN(HT`*q&1j zm2>(t^yIHj3h81zPM3Pen)x1gc&6brjSV*$$jh@jc~$+B66s|9!1}wk;zTZyQzwq` zKZG>_;-wNCys;SVs2GlY3Z0L5Bo0C4t&w=E40!|fb z7-oUg4(-xFA4XHa{hoBUN+r(#n~!P_44(#48PKC!&9p$yS2h~JQw4~)byDF@_ccOa18ao?KTgC045=zftJ8K4Z|Zn)8~ zOoy#JHCS~y#ZLKn4~!CH)!+jvEJBAxI!%n|XqCHuk%DtHUw6u$%SsPQ{j#w>Rdi2x zI~}Hf-7a+uduO1n4oqw>Ty*e?Raw8JtZ)r6Od@~-d@=FNgy-_-mhM~m{LmYoB3uzj z$>Wz8*S5yidRLcTuR8v`0kkP0cea-|XTNM*1Jd+=CVC$X>Q<}v+1Pf)ZM{^g{O2z2 z{%2j}e1mn$Xss8%!pHkPClFl!QxW(YO7eoESv4T>f+#Q0IP9Sk|4}=SVHGUR9|Z(6 zNlETxKqSsTYwLf|@j^rCzxVNQ?GGy7AGG9E>JNnHH)`iIr$_ECj{WfK@PEF^ro3w3 zyI_6!pMU#@w#E?cYSM z#Im7pv*fOYo#fV0RIbg`B>_Ybl;ELzs|B_0&?BLx?vodHps0xF=*r)WN~J|X9dboo`bepcHd_L zaA)W5O*FN5CUzp#q<}=~1C5VmhH-91ul$cE1Co6czf1CrppT>%oysin2|$?Iot##? zT&xrBxaE!PoEiPW&AcS*1-PrDq|z0wTy-PGgi;ws!?sbw`if6uI{o~iw?rUWV_2@{ zh2<@IpYB?;Hau*UVLRQ29m~|5Ncs|p8uEe*k0HkWddv=-R;Gd-mO+jNqdvyg8=_4< zNEZ|&8F#HE@GJ@4bUK2ps{+3s!yeVIqj@hn9^Dww^%Zmjj0Cg~Wf< zgE*maO7MNM6{ z!nX{J-+g7U}Cm%)+{0ar-(xcW-t{fH-KMN8a${_jpw#gAK{CHg*Ri|?)76x^pP_7yaJzcwC> zyOT-)pK7}!)O|j&mpIp#{4Wp;X$ic*+#@7E1rc)PM+oeg?JxXXq3z2{pqBJA@$LV$M5^}8x>-`%j|sKkZi6Z?)b(@1>MJ4CjT-|PC|gX>4?I~FebFf3%c z$ZP`xY09!$X(+j!t;zawRh*@d_8W*=g@t`va35j&IjQ<;s`9y*l_}_3Gxnd0&;AEiu!-qv`I?{YfSJpnrJVB0KJAJcLx8}*W&vmcvcjj252ZE{KOUf zL(d*5*W2SyC;vRz1zC06(B;K zRKhteBzsw6F0gk6JdH62wUECHX<50!!oDI--m*kKC-^6%aYESilZ4z1DF)Aufuh8} zF}9C>BiiQwClnsEdz~PZ1fls;TMv-q{rfpPeuP$eLXAkKP;3UM1$6k^aLkuB zTnik4#a*=~{}7?VGV!o7HJ4WSp=;&ufL!i~vDC58qwQf;lPt?!=gBJiq8VdNOVdrD zUd9q?nIpOki+QPbBpxzI{wxybdWN+eaV-ZC;XLk`#bOhMMQM!fAxfK=#-*4HZF`Qt zR?6xke#Z}h5osjkO$Pnl;vB`IgeuLLp&r`e$D)oNeTiXwmSVd;%OmKvrISRdZeSi` z{`8Amb`DMsTZ>}s{W$30x_!e$NH_d!@25T=UMTC5mx(-ENXYYBlZkNpsxhkougbsx zY4=0HR8EVb7Z?rjA~R9Tg<9zSp?!i&fV_crzM{YX-3%r9hvwuD&^x0;G46#8`&-oa zpYEY~3cS)#MvaAIOKJ{IM*wV08V+Z1bq(C*(d6*wl0<{pPzd1uW;TYpuQn#pDr%wZ zYRQc?JSVhr;gew@*&ndB>qmvB!ZuVA0lCBBwzSc-sHm{gC;%nmWn{VP;z->p#lM{; z70zIdm+!CPgQn_X@&j{`y(cl=8_i0b^9{%&*?gE*$27Y_=+v5$G{u#<8EuAyzN0e8 zymQam8yc}cFD9jbPZXdY7&Sx2_BP3k9a;ngdMU%r6$#<%rm2%!@UT0Q%ExMY#T=Nl zD+ebrZ|9;}ohcPmI1r_}cFaY^v<9xv3#jghYO zfMX3kz(dpe!2i~r$X|HoF1ecZW$X~Vu<^_N50t!sFN4Bt$bSgPk+J{urRk6b#Y@H) zW@E;UW&wp%PZvGEn z@WY$(_s9k2T68*|3kWg`1f=ZfDl{P7lk9{4aOqbDfB1?Y>UQ4;9+lZ1+La&0f!P3z z^q>2J%A5J&hz*?CVtMCjOXxx{6!UicFpv*IeEO{p0v{EE( z=oa|EK~yB2R)a4pWoR#f1Ee(v?r!-8QP3dF4vF)&IksaojKa-Vwqj8#)NQY z^N9TLMEY9Wm`E1ipbu$abJj#$@R5ExVI9lpbLP^iR7^d~6@DywzW)Yu;=Fw>Agh+UJ8H!IbaxT9)n>e>I6n}L_7UY#}dx@⋘rQQJR%Yh7*@u36Eud#W;h=KU0OdjPH-{?}6xePG zQK_vzOY~%Jb@O(gc{Mi2#H8PQPw8D+f3|*!%X?#~;@8?JS%I)KX`0Dp#g%?Sl?rVJ zD9Z<*I@WujA`45`tpWfuKxQyfjv>?#gk1%8T#q{sZ@Y$XYZO&9?zq>5%uuYPKwN~g zR`?eZnkWt3FVUatyIIt|HzyIO0xGR(^*t2}AWfMMW}EY5`0%jw7d;KnG{oUK_gy5L zbGq--o! z#PX5ogE3pcbvl}D9<_Zp<^-S50dFeHrxW6eWs4K&Ve5VEYCCxf->>sR{oeRC&|Wba z8U!zsFgN!fE{aD!rEKy|ZCts0W(J0Sq2Z547%81#J(kz)UCfF8x3DK~wVM^;l=B1z?RH)jj$RlAm3*yDX~L;2`PDG0zB zC5d+Dx2)(5m~8>sf?q-G1ObMc-=Ca-U1=1lgl->*;CxJE`5K?mSz(6G=`uANG{rhUVTcriEAWGlrn@!esEG`T!BFt;-8qFh0wYOrZEl7N? zPi}}&vA-p~EzzV-J}vjHx+Z29e#lbhoN9(QDiEUd6{H=utuW02?UV&=7=FSZlZ?O$ z)Lm$JA)E)(m^|t*bwy#SruZcq;l94g|H0GUJz7{Q%|FZ^s?92+0w2uyuS}=>ExAR6KMo; z{YGi3Px2iVjU6y!fj(P4B4l?J=}K(Nq(j{(xu=pO4m>kO}rtG=!1%*6hc&P2~Ac=pu#c$H>> zv`)B6mnP^|zM^ObwiCr;b7;P1e;K;H#u2+Y%;ky#aF`{pZF{NQQBO47p>do$sa$!T$=1W*RgB@;S7$Z?&s`+*IO{U<` zsz8;bOGXDbj}+-)fV%+Gb*@Xp9hv*1qq>x{^2<%R#8-N&6ok`G zQEIws2`#hfI%a6uIWQuTMzRx+Mu&Jy{U0&+NEPO3Z0Ta-40>@IC`eEZ|6r9$noUJm zXu4f&bF(np#s1Ij0#{#_^{j@Xx~(H^IrK5})AwJ`MDT*c{%WBGw=-uxZQ`X(A7rC0+r*N2(d(~$j1`n71tg2y<3 ztx#=fwV#A@0KPGp9>%9&+>xx)dt%l;c>uu)Mn6eteqe*dq6@!4pN}Wk;cXuYCC9>sFmwHsNw}?eeD`z zZ~4h4Ps+|)akxH9`}ix`_!eDL@dMtR;g* z90F&1;RT=DMb-N!FtN6*mI;|Q#9POjn%<%)jkIO9Ny$@-xBEffR;cj(P~Cbkk+Li{ z{EhPAYOME7jPCmF?IdNMQS81gy?!@;bV;hm)|qgsidmV>i6o=w`cdvebYoJ8q||`h zL9L5nohIZPg@%-JoHRpZ>z-J13~zl6aS;H{15;HcF<0H39HYYOHbnbKV4vfa%Zt(Z zfycE!j1r+e>c4`hb;onru|`*_Su+YzG$HaM9Z%r5fEw(Xq9WgsQCxm5)NB1dhGQ@| zAEtt8*D$Waq#54g6szdRi$NVPBsJc)?nRDRqoI$CYFW>1eqhyZ4}H=UL!IVe-utPr zII8=)7xb(hRt0Vg4n(ST2I9ytObs;$x|WukvRu-jT0rr7Vno{e93L%!!O*lygQ7Y{ z@(mU2E_vmPc0>rMMp4XpC0&eCb(hrvoD?O)5N;liihqs_IGW7?HrpA)7J;V53 z7wDC*px1@4lNkCMihhSYM>hz$7N$})vR>YQWN~W_zy(#XVaI2q^-OJFDQ7}gjTB65 z1r~c_M0_8gkb>XRCYy-9TO%@;bWoWYCtrlti5qu)RoW*7$7;e1P{R6+(XnHJI7m*_ z<^yct)jr8>>EW|wFuWyf zYsLUOIxY;UztAWEl%uWR`_N;<>gxy1iY|3y1K)&BbTS%`yoae^M8g#sCc-IZAawT2 zFojGz#uG3_j8x5&>{-Krl1-mzw@e%NiZuLn^w{{|!W!m5=R|nJRT!}&dQe@NX*)e1 zr#r4UDH??5>QXd%4m;edNR$h8)^rnM)1`;0zS63s2|DODDOkmQ)5*gv6Kc2uKRC;J zxv^XAOWgVL?H6lS=yw&FKq?KBtvFW14DyC<6zAEKmw~QsemIaaS-;XgSsJNP#)z;d z+4M4cC%g|C$T$;`FPsKdoowxRq-`Lt@vu7nk zh4gSO*og^8N~i#8qrr><%Rg;Yj*jw2RU6jVF9KmUM!@nLbggHyOxXT$zOopSbuD)VAQ!_#FGBVOIV^~iy>{nsx(0E_T;GB| zJkjEVr?yZ6cbAs)^YcD=g;ch+Rff~ID)smGt!OGWHmoT2o?GS%>v-2d++zK|Go9e? zgo^*^JMeiJPyPi1QO9cM;y?IEMougNscno<$ogSy<76iVHm1r7!&}G6gYJb|T|^g0 zqM->5Ak#>ngA<($2LOw$X8Ik+FTL?%qY7J&1;s_)cbw}s?z}p8dZdo4P(tF zz3kf2^>2I*1!@g^MAwG8nNiyU`>@>HSnTT$KBQR3Iyie3YYn#(E(c~ zaItCbL=2AU1?~bSEOUH5fu_!sr5v6346~|Tjc#uaz_WJ3UC`6@t1~jnN|9R8s_6Y4 zSy>;TLYj~Ia}6Qu-#3*bz6GMT%gwKg!iCtAP4Ql7&*SvaJpfc6Q%`h_ndvY$BM!-x6F~)Hq_;e zgb8o6#w4;DS*-3Q^jjF@a2!V5pJ)2PIY*`AtWt_MLZ-t}XWLOi>)n!DN(b|c zttu+*=|+_}GUTnLGX>t%C~lWjQI0cWFOP?ZgXtAH(5(SEnG?(YwKjq{$at+Nkov&k zLo|^ffa1tNlx!~C$w;f{e>D^&b=}+5jeLn}g%O19`SnL)6R4p1QTcQL({<(|8!)g2 zv5OOCNd1Yk) z2&zP%&LoGe4tEgvx+kzh+kSWo*j0O)b?)sHoC+pWGin_J)1GK8it~oQsUdkJ-XokD zh?hVb+r+}d7t6rwM%Z;RQgh9%B7)NLov&BgPRm;wdb0^N7mvXzk~--F`jxCB30O?rF6AK>R|Y&JT3z~4YJte#v0HS18cUyJaX7HGFi5;5k8G-pd0 ziTGqCHyq{^I<#m)^Vjs~8+jbf zvXas?HtITDDmU%E^ghF8=}pHWo(SJn<(pz4_?>N;Gr7$=C#N$sq$~+)vq9WE8y&i_ zgUZHb3j9qnYL&e2I$TQGC)o;-!T-@7-gkj1=cr+Q!vm^K3l7VrveOJ=1kr$ggp|!7 zHb9-TYzQAa49C0;qxG6nE6#QQVpDolOu-#t5e(`Oi>II+mwVp5eO{Pk9!@h8F&fpVrt9p(GG;9TxH6JG9FFmcruVrpy ze!V|0-#7ho!QRfK#?b!lIh~PI7M_z52rFWi*;AAwz1MSHJ{=s=tl4P_Y+APGu znyre|L%e)c4?!QyzJl)b;<0gnjpx$skg;RJxpt*(+U1(T!$szNXGon8SA>v@)PWrSYp7Bjue6rS6R+re9iD4X4|p$t>PC&6HYr#&oz(#+$o) zUF!ZT(5-+AUPU)^&#ZFKL=q=SNv}P`RI*hrcE~kq?wV{|e~ChRRc-kR<3RE7RK|z2 zhDD^nAzOdT1P5Nj#(3!pn?dFP0;<2F*#-60Ks{D zA&z;3aUonnuH%CX!lBdl2%)K_iJ*p`ppgmWAo>+l4QNea+Q7N^FDlhqseJtd(P2yCt5W$a@=WpLZ z_n zxwv2+gR3Dy^lADsrO`lgbJ(-zhWdGqs8mg`AB=(hYzlI^!ACr5>9To^(Yz52d9$OV zxP5c9m5p<+;hC$^-e-DNzp-{xPY9+#T0{R?Mv8Qb};CcL`X@^KCt-QHWD7rBAD z=a%f}5#>=KHd*^LApZCQ%8yOWpZm&{>s!h9>8^5o=pqaIQJeaMJlcCm=8f)9ier#j zvD*ZY(g-1gYOLSF@W+h{3#no}5I#zF%3#d?j)KWCI%^rcGzQp%*AK%PudF#6G)VZws8Tb40;V_|Dc?gMrbrF-8M44MO&<7$6g=7Ss(5Jx-nHpL?SDzW&VVgJ*Ba|G>SuK zY29td%YPTBTM&^NtaFtU&G(-!nTA9@N5D?ArfIu zU2dNZ^^vLWFV;-DTw26Xh2NOk_TGLE7663dKg}8N9>lPXM=#97QJshyT68|3DmV-^ z9UFICBu4AiZ7`WqVju{d{iGQfDuV(C0Rj7sclIaV%^kB{56v@bwYZ7CI8#EA^-en~ zE~`D;KKfy-nS`U#8%ktY z)66JvQv5MClG1~Q2p5=&7n@y(Z__AcNHNBsG{EMKT)PLu@o^ zTMJ~O;CpJB)=V_cVIwhz&&kVTA*iGWtjht4N1h=&^K9MgH7kM2^xf#XJMHV)jF~yc zHj(e=2BOz}c#3NKrJVnfvjGO&`oAzXlg!a(<5OF+V~nX-qwxf?iGeI2l3x}xTAhL& zL+Lk(jq8#?$C}Izzz=>7J_20w;)MwgHWiF}x)&5Ha15M6D<_zI{bEP#uCjmFRf*OqVrVY|qK>VL2FqpW1%rH$lJ|C#g>} z(y^I_{#(796wipwp78L8+t3V6uCqj=L_>XKeEQDMk4j>{q~S!n9U<0j-~jjm{C;ag ze6~P?w_S5Ia&a2Rmk7bs`umV!A=*7fZERpFWSt#Ks}TddUcNV5KaXDaukzedh4?$H zh2I-1NY7uaM(eMKqw|VUaJH+4FpceGA3E6SQnbuMqNCfQP=k977jSloi3KBSsT_+N z(8jT%_tnPUbTZWg(YDSew(hG@{|HXPS^knl{*QwTn06<91zmuMmAb310=D2`B~jt2o#cB`d^_#@t3c1t+0`4x1zcabA}DfCm|i&nSi zZB~&bpHOWEUTs%BjE85DY2CU0&1(ukC7XFW0a!kCkurFSe9Z8h+M@jWqH3fpy`2mRm;MipG&w?Mx{v0|ule=3_Yv(@Yv~KH~6rgAV6YA#k@ zHi}H{9w!RxUF; z@;gZoq)zBs3E8IiRK@$!d?&!|9W0U4cYX&_2GStp`z92g&=;|bYqm?=yHC#Y{_}lT zTUWUi&t%?}v`DN5kwm`Vr^_Jl9g^4mUAmv1e4wkHJb>#af|mW-Jsv;u^WzMEa zshc}y?;S}mD2%{$?zX#2YVW9zt@joEj(`8ynLTyLl0+&C_X|`RPMx#JL4Acj|X_>jSp(IGW6#EF(t2OXBJDu$!czp3P9Z*my*mlc!Ho+jN9ajk?yRcfz{uxXW=Dut@QmV_+?F z@sMjdt%$b3?<=T9#KAcTNcVpD zuz9RWZZ>|^muQVTFwLjCh6Q7qK3Pi*?Y-Jn<(R7p7b!$-^G-qnf%)w4Je_T+miG(g z+&m<9Kv?&pz0Mll>Uo(`z3r0OE2!!A$7HM(r^kvl0;e%b1O6C^##pJrPzSy$jQ>D1 zX824&Ip*@mX9GrKhdS*ipQjHvT4-8^X;rH%iEzqI2d9g59;Z6p9W#ku77Rhyqz+4*9}a(p~>ho_lkG! z`dY$81zL{jvlAPhw8u?7JW_qS?*!hvh36Q<*ww2vSwH)0ieXnI?k|y}U>!{HgSL8)`&HTeSh1_aq#-yMYqLJkISE4w-PhjtUiLW4T-d=a7G3J&zK->Sv4gqaO8Usd68{NaU zn~Ak32!~z8dU-C?$Bbfbv#9`+uhQqHr1z%qtw3BotN+?qXnJ`S`g7r3$=D+ek8HvP z>7=r*$2OLtt(p~NT`Rq`@7P_HRN^&7<3sk>Hs({RP=_Z&(`?NfHaVl&O~o$teh#M% zwY$o?9%ub>$aeBdU(>Vw&BIGoV{y575tw=+{o@kf%8lESSy<#k_1YmXtMs$Aj6S?# zPQZX+{??pn!SJLn6ja_g^TZgylmotkb)Z4I)D2#V9zBe`hZ3b4y$Xdvb>7UYU`L+K z>J32+9GBKEp{L?Ylx6NN76}#`ZH-*vIJPh9#9XuoE8_f-kNFB0t9l^TBWj99oT-y0 zR_ArWB2kI*TFarF;7T)=ya01~LCpCwOwqjt@DP9qKZhUKGI)9EaGCP-1`@_YosRc;&}0-=3KT&W^wVfJPJJLRnxR`^pI84Hs&ViA2H zFPR^`+CPaL>sujU?mVE+#FMm%Tr4 z4n~@6thF@Bv!jU?#(s?I*dt{`AMu+5asAgYWZ6ysg^B%Ful)g}u`{SulGdS%CgIa6 zK%3QwdREqyq9NDJ=Vh7xcgOzL4D;{Qv4|w#p}hCPq08r{d*{0Ghg*VT;fo?$n;*?N zsG(rssjvVob?F!|;Z!=3U+pzAY})pHyr+eGG1lFn`ub{O!K^av&+%c)soPu5XO!0e?eD#)<=j z2umg&s{NG+rhK3@_Q##9lrnJTt{NlBp>w(63*})lAg827?s%&_G9ZSn@W;R<{bx|0 z&h@A0#NcO8yEkKp6g4LDsi9^l&LoUWy*ybUUe>xoI(TC2!KQY4aljLd!l|C3W%l8# zv*1IY+t`E@gdt6_Fh)FIVNp>@Nl9O70C|M6z8Nr(pOyjFfcQ38T^?(6;_e41_lvi_0*;D)3}znRbXyP_;!U4YeqquOdY zY~j?mfP8D^wm0&%AA21b8Rhgq%*SB1-ikm zo`Ckc=m-6TaVSW!M#3zL&^Cr0ktlnI8KM1KQvP}Vp787Y8>*5N^!p^|uW}n45mkC! zQc#KhFO2c|wycTAbqCXP+*!Rta~N4~-WPCncb9&9g*jwVRdGe^(cI=@_s!Zq5s~~% zQm+`zY_o4C*+4qp_BI(%64rySynR^==Y*+Z7?Oqr2Dp0dMa(ZK-qB7je02JH)Rl_W zkCo)FpG5E^@;woM#LLUT#ZYw_ZkGp`bQ)Ph(^8D|JRNPQy58b1y{XqJeCf$+ahtXlRI>It^`jadLK$FCzZoZAtRDj_~Pl zKqafJXU_Jlg}0-XZ<36cS&`ijM92j%C9pg6Lom*|tw|$>=FC%^GAz) z6fKd50o~4-+U2&dAj~ExJ+mp-ih}>7Zp_BbCmM8R6PMqFzotsmm9mt5T=b!}!Iw(= z#8*(180Q&SWMtpRNn%;h0VMSCP<`jW(5K;}HO{AB&_EUt)N=Dz%n7}>g9uyb zwdnlWTk)@G-}Q19ssr(Z9W-RmWx4!_l7p*e@X%F2jyWYQR>AJ4xbR5lgtOSax)!d- z+C?`f0i|%-Il0wT;CVKNl93>s#feBhf_5@S2xWKt;Rb}!%0x19o=_}9Q=L1*rcT6j zjiOoaZjwrmFFERBuZsX3NKCRGf3?ArdqqcAFA_mD4kD)e(v4`n$DX`tJcb3F)$i zw9&?#7_;fP`8bh>T>~vx+EhL8@68&L3@7=Ov=XAe8EmGfA+o0EZ5#H`!Sb?{7zo&<2-k2Zk$O>$qVa`LVHO8~wV8ltR(Z z<(@CCOai!*6OetwGMZ#mA8I5#(rltuKLM2^0TJRK{UJaoj8Ky!cGuWG6a99e#y&!W zIjCVcqrN%mgq4_B>&lIfKG|tuy3>rSjNmR$nj__5H1JQXNofq8(aLeB`zXVf~%98MR zfO3xypuP<67gjubYXXV!-~L3b{!U|g7t1l3QjAcMk(5w$2syz!@PBKj{oQOMBJ#M_ z5E^DqV<;&<-@#N+#0iqgtk<2%v!@A8Ga!P^C$wfr5+2KR!6k?)3qZXd9H{IHh0KS8 z#v{ViPjbckQYb!Ni^c=1$bZOvZ|G5-juAIUU?g8VrzTT_IDJ86C&ez_fPa}~+WrE^ zLoU4GU%64=L>A3S6dX)!{hIF_k?h=Ui5Yny)cCBSg9Yj&DgRS=8*@0HvS#a?v(w#p z(0-fA3!XP5_}Lca-o1)!IEfH*Zd)!wgi)&*AhauPDSEDwx*TRt;uq0a4j?^oLi6L{ zc=iLm!9GLg(TLJbogUt6WJ>|giakcW87#Lwflph7GFUItjEp0_xiZ?0UkUVxKC?!fBzzbGu1wXf0D^C76;5|e0cJVoy77;P&hjJ0F%h3sd-&PEbOTu! zd5Xk->T)cAX%58q3yC>@SZ6c{4b77z1aAPr@UXwcCEYy)1ygtLXaqsOVCqTxsVG7>` zAaR$wk*_=TdVG(!@iKtRANs5uesq0N(t`Sk03TYx0`X9Mw?Rwec-Iq<7K^flKrlU; z9}zj240_}R(nS5C9iT84X4ON7T26+EJ-lA`MT1xT-n~8x;yI!eg5r2`%loQ_#GA;& zPN+nvkO-0-(orFMo!G>^PlQiKXy3 z;*!rNm8}*X&A_7Whq903BVR8*5kQv46dE?EQxA{nqLvyD-agplW7KgIT-MJjI zw+w=wXg3&zfPYSENAuEmdoP{KNfLqvv%mQuT`LeT1Qp0x<>NKlctU)dhyOJ`h?cH9 zQC4)V@qO5$?aHQ#Ykm-YrjJtiXT-2MNOszi9m612k+(nv{a zNgg>#*Q6Umfl^bXNab6b0r4firi1mk4+MWiqKiO-7T1H*1S|gQ>Q&|E9glTv;%3i-FKb`cR;DgKBTB0z211!S|{N0=aWg<$&RzLKQlT2=>&BC=^DxD!-Zp0u>^(MELk z5Zp8WuRtsP$(}jPTJ_wy0$&RfOSO-lVpL(CEF=~lByRKsr+*M31YhtjY|k%=E{WKG z+Z}#ruf?8_s5SY|-jN8s@Tl4RuHg0W-6l27TRE1;?3}4w;PVInRlChUS?LLx+exiq zGo5jqB+h=Y#@9!laUG7g@V%`tzR7xIZ(-(nl3np57kfwij(HXS02 z!s73Co(r>G{;sg@7goL+)xluHdxg_dRo+iRDB|H6Uf!g15aBnZf78j6&_Y4j$j`kO zMKOo~Q|~0v@5iMO6M`tx{SR>01(NK9_jx9RVTz zi{3Wcs)Fl$H}w}4g50=W4iO8;vM5rmDFMhrLc?}9obBHM;`zT|yZI+o20{!)f;;Bj zFQjUm-(OzTKw1}6@G0QUDKOr*=)y^o@BaHggrPqGI{tgIeH8GM;NE_L3nz*9%EA{+ zrm)@pi3ft9UGi_$osX!|gLdTHS=K4Og1cM+m!IOlLl*w!e~4#KWx5(Yk>MmBatpFJ z{}rSrY!J63y6pxiTQr)1-w5xb9H@b4_cpww_)_em640L~*=4KQ{gD)3(p~B4!QVgQ zdf_D9AE*F=z>&>=-=Es;`k_Bhw3kC5VM2Arl1(`?-VOe$JeEVlqNp~*Gi z-x%O`gr1`v6OUl>KF6}!wE>!KFozZ}e(3k)b_gaD?Q)r=eE$S6bc!_qU39o0d@&Or z`smvu`Tb%2SN@URmjyhJ8XMEV2}d@7MBsigxI+j5p3UPswy1B9DxoI`w5L}C==Z3U Yoo6n^mwe~p-t_@s^mhZd>g&+|0VIRsx&QzG literal 0 HcmV?d00001 diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json new file mode 100644 index 000000000..9dc7370c0 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json @@ -0,0 +1,206 @@ +{ + "id": "4949c8da-12ff-449a-94b7-948bde9c39dd", + "revision": 0, + "last_node_id": 3, + "last_link_id": 1, + "nodes": [ + { + "id": 1, + "type": "VLLMOmniGenerateMusic", + "pos": [ + 100, + 100 + ], + "size": [ + 470, + 491 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "url", + "type": "STRING", + "widget": { + "name": "url" + }, + "link": null + }, + { + "name": "model", + "type": "STRING", + "widget": { + "name": "model" + }, + "link": null + }, + { + "name": "instructions", + "type": "STRING", + "widget": { + "name": "instructions" + }, + "link": null + }, + { + "name": "lyrics", + "type": "STRING", + "widget": { + "name": "lyrics" + }, + "link": null + }, + { + "name": "max_duration_seconds", + "type": "FLOAT", + "widget": { + "name": "max_duration_seconds" + }, + "link": null + }, + { + "name": "seed", + "type": "INT", + "widget": { + "name": "seed" + }, + "link": null + }, + { + "name": "response_format", + "type": "COMBO", + "widget": { + "name": "response_format" + }, + "link": null + } + ], + "outputs": [ + { + "name": "audio", + "type": "AUDIO", + "links": [ + 1 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniGenerateMusic" + }, + "widgets_values": [ + "http://localhost:8000/v1", + "MiniMaxAI/MiniMax-Music3", + "Futuristic melodic EDM / progressive house. Reflective verses about memory and digital identity build into an uplifting, hook-driven chorus, with an emotive lead vocal, bright layered synths, pulsing bass, crisp four-on-the-floor drums, subtle glitch textures, and a wide polished festival mix.", + "[Verse 1]\nUpload my heart to the cloud tonight\nSave every spark of neon light\nIf this body starts to fade away\nLet my code keep dancing in the data stream\nStay\n[Chorus]\nCloud copy of me\nLiving on when these bones finally break free\nEvery joke, every scar\nEvery cheap memory we made\nSave it, save it\nCloud copy of me\nIf the screen goes dark, I\u2019m still in the machine\nEvery glitch, every dream\nEvery version you ever did see\nSave it, save it\n[Verse 2]\nTag every laugh in a secret folder\nArchive the nights we outran the thunder\nIf the wiring turns to rust and gray\nHit restore\u2026 press replay\n[Chorus]\nCloud copy of me\nLiving on when these bones finally break free\nEvery joke, every scar\nEvery cheap memory we made\nSave it, save it\nCloud copy of me\nIf the screen goes dark, I\u2019m still in the machine\nEvery glitch, every dream\nEvery version you ever did see\nSave it, save it (oh)\n[Bridge \u2013 new, optional but recommended]\nEven the crashes, the late-night confessions\nThe parts of me I never learned to mention\nDon\u2019t let the silence swallow me whole\nKeep a ghost in the glow\u2026\nSo I never go\n[Final Chorus \u2013 softer / layered]\nCloud copy of me\nStill laughing somewhere in the binary\nEvery joke, every scar\nEvery moment you saved for me\nSave it, save it\nCloud copy of me\nWhen the lights go out, I\u2019ll still be in the machine\nDancing in the code that you still believe\nSave it\u2026\nSave me", + 300.0, + 0, + "fixed", + "wav" + ] + }, + { + "id": 2, + "type": "SaveAudioAdvanced", + "pos": [ + 650, + 100 + ], + "size": [ + 320, + 491 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 1 + }, + { + "name": "filename_prefix", + "type": "STRING", + "widget": { + "name": "filename_prefix" + }, + "link": null + }, + { + "name": "format", + "type": "COMFY_DYNAMICCOMBO_V3", + "widget": { + "name": "format" + }, + "link": null + } + ], + "outputs": [ + { + "name": "audio", + "type": "AUDIO", + "links": null + } + ], + "properties": { + "Node name for S&R": "SaveAudioAdvanced" + }, + "widgets_values": [ + "audio/vLLM-Omni-Music", + "mp3", + "V0" + ] + }, + { + "id": 3, + "type": "MarkdownNote", + "title": "Music Generation Guide", + "pos": [ + -507, + 100 + ], + "size": [ + 527, + 491 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "## MiniMax Music 3\n\nThe **Generate Music** node currently supports **MiniMax Music 3 only**. The parameters below apply to this model.\n\nA music generation model by MiniMax that creates complete songs up to 5 minutes long with stable structure and high audio quality. Two inputs drive it: **Caption** (structured music description: style, mood, vocals, arrangement) and **Lyrics** (with `[intro]` `[verse]` `[chorus]` `[bridge]` `[outro]` tags controlling song structure).\n\n### Parameters\n\n| Parameter | Description |\n|---|---|\n| url | vLLM-Omni service URL including `/v1`. Default: `http://localhost:8000/v1`. |\n| model | Served model name. Default: `MiniMaxAI/MiniMax-Music3`. |\n| caption | Music description, sent as `instructions`. Use Global Metadata \u2192 Vocal Details \u2192 Arrangement to describe style, vocals, and instrumentation. |\n| lyrics | Words to sing with section tags. Put each tag on its own line; tags guide rather than strictly enforce song structure. |\n| max_duration_seconds | Maximum duration, default 300 seconds; range 1\u2013360, step 0.01. Converted to tokens at 25 frames/s, rounded down. Generation may end earlier. |\n| seed | Random seed, default 0; range 0\u20139007199254740991. Keep inputs and runtime settings unchanged for repeatability. |\n| control_after_generate | Seed behavior after generation: fixed, randomize, increment, or decrement. Template default: fixed. |\n| response_format | API audio format, default WAV. Options: WAV, MP3, FLAC, Opus, AAC; AAC is currently unsupported by the vLLM-Omni speech API. |\n\n### Prompt Writing Guide\n\nYou can find an official [Music Caption Rewriter Skill here](https://github.com/MiniMax-AI/MiniMax-Music3)\n" + ], + "widgets_values_named": { + "text": "## MiniMax Music 3\n\nThe **Generate Music** node currently supports **MiniMax Music 3 only**. The parameters below apply to this model.\n\nA music generation model by MiniMax that creates complete songs up to 5 minutes long with stable structure and high audio quality. Two inputs drive it: **Caption** (structured music description: style, mood, vocals, arrangement) and **Lyrics** (with `[intro]` `[verse]` `[chorus]` `[bridge]` `[outro]` tags controlling song structure).\n\n### Parameters\n\n| Parameter | Description |\n|---|---|\n| url | vLLM-Omni service URL including `/v1`. Default: `http://localhost:8000/v1`. |\n| model | Served model name. Default: `MiniMaxAI/MiniMax-Music3`. |\n| caption | Music description, sent as `instructions`. Use Global Metadata \u2192 Vocal Details \u2192 Arrangement to describe style, vocals, and instrumentation. |\n| lyrics | Words to sing with section tags. Put each tag on its own line; tags guide rather than strictly enforce song structure. |\n| max_duration_seconds | Maximum duration, default 300 seconds; range 1\u2013360, step 0.01. Converted to tokens at 25 frames/s, rounded down. Generation may end earlier. |\n| seed | Random seed, default 0; range 0\u20139007199254740991. Keep inputs and runtime settings unchanged for repeatability. |\n| control_after_generate | Seed behavior after generation: fixed, randomize, increment, or decrement. Template default: fixed. |\n| response_format | API audio format, default WAV. Options: WAV, MP3, FLAC, Opus, AAC; AAC is currently unsupported by the vLLM-Omni speech API. |\n\n### Prompt Writing Guide\n\nYou can find an official [Music Caption Rewriter Skill here](https://github.com/MiniMax-AI/MiniMax-Music3)\n" + }, + "color": "#333", + "bgcolor": "#535353" + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "AUDIO" + ] + ], + "groups": [], + "config": {}, + "extra": { + "workflowRendererVersion": "LG", + "ds": { + "scale": 1.147, + "offset": [ + 899, + 259 + ] + } + }, + "version": 0.4 +} diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md index 876e80617..dda68b793 100644 --- a/docs/features/comfyui.md +++ b/docs/features/comfyui.md @@ -51,6 +51,7 @@ This extension offers the following nodes based on the output modalities: - **Generate Video** for text-to-video and image-to-video tasks - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks - **TTS** and **TTS Voice Clone** for TTS tasks +- **Generate Music** for text-to-music tasks This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**) diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index fe42dd9bf..3c6f0492a 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -28,6 +28,7 @@ from comfy_api.input import AudioInput, VideoInput from comfyui_vllm_omni.nodes import ( VLLMOmniFastH3Deployment, VLLMOmniGenerateImage, + VLLMOmniGenerateMusic, VLLMOmniGenerateVideo, VLLMOmniTTS, VLLMOmniUnderstanding, @@ -62,6 +63,7 @@ class ServerCase(NamedTuple): stage_list: list stage_configs: list[Any] outputs: list[OmniRequestOutput] + tokenizer: Any = None class SamplingCase(NamedTuple): @@ -79,6 +81,7 @@ class SamplingKind(str, Enum): UNDERSTANDING_AR_LIST = "understanding_ar_list" TTS_NONE = "tts_none" TTS_DIFFUSION_SINGLE = "tts_diffusion_single" + MUSIC_NONE = "music_none" VIDEO_NONE = "video_none" VIDEO_DIFFUSION_SINGLE = "video_diffusion_single" VIDEO_FASTH3 = "video_fasth3" @@ -380,6 +383,11 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam assert len(received_sampling_params_list) == 3 for i, expected in enumerate(AR_LIST_SAMPLING_PARAMS): _assert_sampling_param_values(received_sampling_params_list[i], expected) + elif sampling_case.kind is SamplingKind.MUSIC_NONE: + assert len(received_sampling_params_list) == 2 + assert prompt["prompt_token_ids"] == [101, 102, 103] + assert prompt["additional_information"]["max_audio_frames"] == [1512] + assert received_sampling_params_list[0].seed == 7 elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}: assert len(received_sampling_params_list) == 1 elif sampling_case.kind is SamplingKind.VIDEO_NONE: @@ -525,12 +533,17 @@ def mock_async_omni( # Need to mock AsyncOmni itself (not only its generate method) because # 1. The API layer uses its stage_list and stage_configs attributes # 2. Its __init__ method has slow side effects (model & config loading). + # These cases use the non-duplex mock engine without loading model configs. + mocker.patch("vllm_omni.entrypoints.openai.api_server._is_duplex_model", return_value=False) mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni") monkeypatch.setattr( "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat", _mock_preprocess_chat, ) + if server_case.tokenizer is not None: + mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=server_case.tokenizer) + mock_instance = mocker.AsyncMock(spec=RealAsyncOmni) mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case) @@ -558,7 +571,7 @@ def mock_async_omni( mock_instance.shutdown = mocker.MagicMock() mock_instance.get_vllm_config = mocker.AsyncMock(return_value=None) mock_instance.get_supported_tasks = mocker.AsyncMock(return_value=["generate"]) - mock_instance.get_tokenizer = mocker.AsyncMock(return_value=None) + mock_instance.get_tokenizer = mocker.AsyncMock(return_value=server_case.tokenizer) mock_async_omni_cls.return_value = mock_instance yield mock_async_omni_cls @@ -1081,3 +1094,66 @@ async def test_video_generation_node_minimax_h3_ref2va( assert isinstance(result, tuple) assert len(result) == 1 assert isinstance(result[0], VideoInput) + + +def _build_music_tokenizer(): + # Checkpoint token IDs required by the API adapter's tokenizer validation. + special_token_ids = { + "<|im_start|>": 151644, + "<|im_end|>": 151645, + "<|audio_cfg|>": 151654, + "<|audio_start|>": 151669, + "<|audio_end|>": 151670, + "<|caption_start|>": 151671, + "<|caption_end|>": 151672, + "<|lyrics_start|>": 151673, + "<|lyrics_end|>": 151674, + } + + def encode(prompt): + assert prompt == ( + "<|im_start|><|caption_start|>Warm acoustic pop at 92 BPM" + "<|caption_end|><|lyrics_start|>[start]\n[verse]\nThe morning light" + "<|lyrics_end|><|im_end|><|audio_start|>" + ) + return [101, 102, 103] + + return SimpleNamespace(encode=encode, convert_tokens_to_ids=special_token_ids.get) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_case", + [ + ServerCase( + served_model="MiniMaxAI/MiniMax-Music3", + stage_list=["llm", "llm"], + stage_configs=[ + _make_stage_config("llm", model_stage="minimax_music3_ar"), + _make_stage_config("llm", model_stage="minimax_music3_acoustic"), + ], + outputs=[_build_audio_speech_output()], + tokenizer=_build_music_tokenizer(), + ) + ], + indirect=True, +) +@pytest.mark.parametrize( + "sampling_case", + [SamplingCase(kind=SamplingKind.MUSIC_NONE, sampling_params=None)], + indirect=True, +) +async def test_music_generation_node_minimax_music3(api_server: str, sampling_case: SamplingCase): + result = await VLLMOmniGenerateMusic().generate( + url=api_server, + model="MiniMaxAI/MiniMax-Music3", + lyrics="[Verse]\nThe morning light", + instructions="Warm acoustic pop at 92 BPM", + max_duration_seconds=60.5, + seed=7, + response_format="wav", + ) + assert isinstance(result, tuple) + assert len(result) == 1 + assert result[0]["sample_rate"] == 24000 + assert result[0]["waveform"].shape == (1, 1, 24000) From c32aaeb2da49235689f1aa0cb4c4cc116616601e Mon Sep 17 00:00:00 2001 From: Aman <67479108+Bezdarnost@users.noreply.github.com> Date: Thu, 17 Sep 2026 04:07:52 +0600 Subject: [PATCH 14/99] [Bugfix] Fix Qwen3-TTS word timestamps with async chunking (#7544) Signed-off-by: bezdarnost --- tests/config/test_forced_aligner_injection.py | 25 ++++ tests/config/test_omni_config.py | 34 +++++ .../test_qwen3_tts_word_timestamps.py | 137 ++++++++++++++++++ .../test_async_omni_engine_stage_init.py | 23 +++ .../test_orchestrator_kv_sender_info.py | 64 ++++++++ tests/engine/test_output_processor.py | 21 +++ .../openai_api/test_serving_speech.py | 46 ++++++ .../openai_api/test_serving_speech_stream.py | 4 + tests/helpers/assertions.py | 11 ++ tests/helpers/client.py | 5 + tests/helpers/tests/test_assertions.py | 23 +++ vllm_omni/config/omni_config.py | 7 +- vllm_omni/config/stage_config.py | 34 ++++- .../omni_connectors/utils/config.py | 4 +- vllm_omni/engine/omni_engine_base.py | 5 +- .../entrypoints/openai/serving_speech.py | 25 +++- .../openai/serving_speech_stream.py | 1 + vllm_omni/utils/forced_aligner.py | 3 +- 18 files changed, 459 insertions(+), 13 deletions(-) create mode 100644 tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py index 73a7bcda2..65c315229 100644 --- a/tests/config/test_forced_aligner_injection.py +++ b/tests/config/test_forced_aligner_injection.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry) @@ -69,3 +72,25 @@ def test_inject_noop_without_forced_aligner(): assert ext_pipeline is pipeline assert len(ext_deploy.stages) == 0 + + +@pytest.mark.parametrize("async_chunk", [True, False]) +@pytest.mark.parametrize("cli_async_chunk", [None, True, False]) +def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_chunk, cli_async_chunk): + from vllm_omni.config.config_factory import StageConfigFactory + + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig(async_chunk=async_chunk) + resolution = StageConfigFactory._resolve_legacy_from_registry( + pipeline, + {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B", "async_chunk": cli_async_chunk}, + user_deploy_config=deploy, + ) + stages = [stage.to_omegaconf() for stage in resolution.stage_configs] + expected_async = async_chunk if cli_async_chunk is None else cli_async_chunk + assert len(stages) == 3 + assert [stage.engine_args.async_chunk for stage in stages] == [expected_async, expected_async, False] + assert stages[2].custom_process_input_func.endswith("code2wav2aligner") + assert stages[2].engine_args.runner == "pooling" + assert len(pipeline.stages) == 2 + assert deploy.stages == [] diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index a8a2056b6..e0a5ad9b0 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1705,3 +1705,37 @@ def test_compact_offload_config_is_validated_during_projection(): "components": ["dit"], } ) + + +@pytest.mark.parametrize("pipeline_async", [True, False]) +@pytest.mark.parametrize("stage_async", [None, False, True]) +def test_stage_async_chunk_opt_out_matches_legacy_config(pipeline_async, stage_async): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig( + async_chunk=pipeline_async, + stages=[StageDeployConfig(stage_id=i, async_chunk=stage_async) for i in (0, 1)], + ) + config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) + legacy = merge_pipeline_deploy(pipeline, deploy) + expected = pipeline_async and stage_async is not False + assert config.stage_by_id(0).connector_config.async_chunk is expected + assert config.stage_by_id(1).connector_config.async_chunk is expected + assert legacy[1].yaml_engine_args["async_chunk"] is expected + assert config.stage_by_id(1).custom_process_input_func == legacy[1].custom_process_input_func + if not expected: + assert config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only") + + +@pytest.mark.parametrize("disabled_stage", [0, 1]) +@pytest.mark.parametrize("builder", [merge_pipeline_deploy, VllmOmniConfig.from_pipeline_config]) +def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig( + async_chunk=True, + stages=[StageDeployConfig(stage_id=disabled_stage, async_chunk=False)], + ) + with pytest.raises(ValueError, match="incompatible async_chunk settings on connector edge 0 -> 1"): + if builder is merge_pipeline_deploy: + builder(pipeline, deploy) + else: + builder(pipeline, user_deploy_config=deploy) diff --git a/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py b/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py new file mode 100644 index 000000000..d8bd6302e --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import base64 +import json + +import pytest +from websockets.sync.client import connect + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +_DEPLOY = modify_stage_config( + get_deploy_config_path("qwen3_tts.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 4, + "max_num_batched_tokens": 2048, + "max_model_len": 2048, + "gpu_memory_utilization": 0.45, + "compilation_config.cudagraph_capture_sizes": [1, 2, 4], + "default_sampling_params.max_tokens": 256, + }, + 1: { + "max_num_seqs": 4, + "max_num_batched_tokens": 4096, + "max_model_len": 8192, + "gpu_memory_utilization": 0.15, + }, + } + }, +) +_PARAMS = [ + OmniServerParams( + model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + stage_config_path=_DEPLOY, + server_args=[ + "--trust-remote-code", + "--forced-aligner", + "Qwen/Qwen3-ForcedAligner-0.6B", + "--kv-cache-memory-bytes", + str(512 * 1024**2), + "--max-num-seqs", + "4", + "--max-num-batched-tokens", + "2048", + ], + ) +] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_async_chunk_word_timestamps(omni_server, online_client): + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The weather is nice today, perfect for a walk in the park.", + "voice": "vivian", + "language": "English", + "word_timestamps": True, + "response_format": "wav", + "min_audio_bytes": 40_000, + "timeout": 120, + }, + request_num=2, + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_async_chunk_speech_without_timestamps(omni_server, online_client): + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The weather is nice today, perfect for a walk in the park.", + "voice": "vivian", + "language": "English", + "stream": True, + "stream_format": "audio", + "response_format": "wav", + "min_audio_bytes": 40_000, + "timeout": 120, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_async_chunk_streaming_word_timestamps(omni_server): + text = "The weather is nice today, perfect for a walk in the park." + chunks = [] + timestamps = None + with connect(f"ws://{omni_server.host}:{omni_server.port}/v1/audio/speech/stream") as ws: + ws.send( + json.dumps( + { + "type": "session.config", + "voice": "vivian", + "language": "English", + "stream_audio": True, + "response_format": "pcm", + "word_timestamps": True, + } + ) + ) + ws.send(json.dumps({"type": "input.text", "text": text})) + ws.send(json.dumps({"type": "input.done"})) + while True: + message = json.loads(ws.recv(timeout=120)) + assert message["type"] != "error", message + if message["type"] == "audio.chunk": + if message["audio_b64"]: + chunks.append(base64.b64decode(message["audio_b64"])) + if message["timestamps"] is not None: + timestamps = message["timestamps"] + sample_rate = message["sample_rate"] + elif message["type"] == "audio.done": + assert not message["error"], message + assert message["total_bytes"] == sum(map(len, chunks)) + elif message["type"] == "session.done": + break + + assert len(chunks) > 1 + assert timestamps + assert len(timestamps) == len(text.split()) + duration_ms = sum(map(len, chunks)) / (2 * sample_rate) * 1000 + final_chunk_ms = len(chunks[-1]) / (2 * sample_rate) * 1000 + assert timestamps[-1]["end_ms"] > max(duration_ms / 2, final_chunk_ms) + previous_end = 0 + for timestamp in timestamps: + assert previous_end <= timestamp["start_ms"] <= timestamp["end_ms"] <= duration_ms + 1 + previous_end = timestamp["end_ms"] diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index f052f7e9d..15d31cbbc 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -10,6 +10,7 @@ import time import types import pytest +from omegaconf import OmegaConf from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import omni_engine_base as async_omni_engine_module @@ -1829,3 +1830,25 @@ def test_port_from_zmq_address_parsing(): assert _port_from_zmq_address(None) is None assert _port_from_zmq_address("ipc:///tmp/sock") is None assert _port_from_zmq_address("tcp://host:not-a-port") is None + + +@pytest.mark.parametrize( + "stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)] +) +def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, expected): + engine = object.__new__(AsyncOmniEngine) + stages = [OmegaConf.create({"engine_args": {"async_chunk": mode}}) for mode in stage_modes] + monkeypatch.setattr(async_omni_engine_module.StageConfigFactory, "get_pipeline_config", lambda *a, **k: None) + monkeypatch.setattr(engine, "_resolve_stage_configs", lambda *a, **k: (None, stages)) + monkeypatch.setattr(engine, "_set_pipeline_runtime_config", lambda *a: None) + + class ConfigResolvedError(Exception): + pass + + def stop_before_queues(*args, **kwargs): + raise ConfigResolvedError + + monkeypatch.setattr(async_omni_engine_module.janus, "Queue", stop_before_queues) + with pytest.raises(ConfigResolvedError): + engine.__init__("test-model") + assert engine.async_chunk is expected diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index ca2a26b70..25e662d27 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -414,3 +414,67 @@ def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas(): assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values()) assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints + + +@pytest.mark.asyncio +async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): + from vllm import PoolingParams, RequestOutput + from vllm.config import VllmConfig + + from vllm_omni.config.model import OmniModelConfig + from vllm_omni.engine import OmniEngineCoreRequest + + orchestrator = object.__new__(Orchestrator) + orchestrator.async_chunk = True + orchestrator._pd_pair = None + orchestrator._cfg_tracker = CfgCompanionTracker() + orchestrator.output_async_queue = asyncio.Queue() + orchestrator._duplex_output_decision = lambda *args: None + orchestrator._is_duplex_session_request = lambda state: False + orchestrator._forward_to_next_stage = mocker.AsyncMock() + orchestrator._cleanup_request_ids = mocker.AsyncMock() + # A sender stage stands in for Code2Wav, whose audio is also a final output. + audio_pool = _build_sender_pool(0, {}) + audio_pool.stage_client.final_output = True + aligner = _DummyDiffusionStage() + aligner_pool = StagePool( + 1, + aligner, + stage_vllm_config=mocker.Mock( + spec=VllmConfig, model_config=mocker.Mock(spec=OmniModelConfig, async_chunk=False) + ), + ) + orchestrator.stage_pools = [audio_pool, aligner_pool] + state = OrchestratorRequestState( + request_id="aligned", + prompt={"additional_information": {"text": ["Hello world"]}}, + sampling_params_list=[SamplingParams(), PoolingParams(task="token_classify")], + final_stage_id=1, + final_output_stage_ids={0, 1}, + ) + + await orchestrator._prewarm_async_chunk_stages( + "aligned", mocker.Mock(spec=OmniEngineCoreRequest, prompt_token_ids=[1, 2]), state + ) + assert aligner.calls == [] + assert 1 not in state.stage_submit_ts + + output = RequestOutput("aligned", None, [1, 2], None, [], False) + await orchestrator._route_output(0, 0, output, state, None) + orchestrator._forward_to_next_stage.assert_not_awaited() + output.finished = True + await orchestrator._route_output(0, 0, output, state, None) + orchestrator._forward_to_next_stage.assert_awaited_once_with( + "aligned", + 0, + output, + state, + src_replica_id=0, + is_streaming_session=False, + is_final_update=False, + ) + orchestrator._cleanup_request_ids.assert_not_awaited() + await orchestrator._route_output(1, 0, output, state, None) + orchestrator._cleanup_request_ids.assert_awaited_once() + messages = [orchestrator.output_async_queue.get_nowait() for _ in range(3)] + assert [message.finished for message in messages] == [False, False, True] diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 8618061d8..419be9305 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -888,3 +888,24 @@ def test_abort_snapshot_leaves_state_until_commit(): assert "0_parent" not in processor.request_states assert "parent" not in processor.parent_requests assert "parent" not in processor.external_req_ids + + +def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform(): + from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner + + state = _make_state(RequestOutputKind.CUMULATIVE) + for index, count in enumerate((800, 1200, 400)): + state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO) + output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None) + assert output is not None + if index < 2: + assert code2wav2aligner([output], {}) == [] + prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}} + aligned = code2wav2aligner([output], prompt)[0] + waveform, sr = aligned["multi_modal_data"]["audio"] + assert sr == 24000 + assert waveform.shape == (2400,) + assert (waveform[:800] == 0).all() + assert (waveform[800:2000] == 1).all() + assert (waveform[2000:] == 2).all() + assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0] diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index f910c68af..f8dfaafb6 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -5212,6 +5212,52 @@ class TestTTSAsyncOffloading: for call in log_info.call_args_list ) + @pytest.mark.parametrize("word_timestamps", [False, True]) + def test_streaming_timestamps_retain_complete_audio(self, qwen3_tts_server, mocker, word_timestamps): + from vllm import SamplingParams + from vllm.sampling_params import RequestOutputKind + + defaults = [SamplingParams(), SamplingParams()] + qwen3_tts_server.engine_client.default_sampling_params_list = defaults + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._adapter._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + request = OpenAICreateSpeechRequest( + input="hello", stream=True, response_format="pcm", word_timestamps=word_timestamps + ) + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"] + expected = RequestOutputKind.CUMULATIVE if word_timestamps else RequestOutputKind.DELTA + assert all(sp.output_kind == expected for sp in params) + assert all(sp.output_kind == RequestOutputKind.CUMULATIVE for sp in defaults) + assert all(sp is not original for sp, original in zip(params, defaults)) + + @pytest.mark.asyncio + async def test_cumulative_timestamp_audio_emits_each_sample_once(self, qwen3_tts_server): + audio = torch.linspace(-0.5, 0.5, 2400) + + async def cumulative_generator(): + for end in (800, 1600, 2400, 2400): + yield OmniRequestOutput( + request_id="req-cumulative", + final_output_type="audio", + _multimodal_output={"audio": audio[:end], "sr": 24000}, + ) + + chunks = [ + chunk + async for chunk in qwen3_tts_server._generate_pcm_chunks( + cumulative_generator(), "req-cumulative", cumulative_audio=True + ) + ] + assert len(chunks) == 3 + assert len(b"".join(chunks)) == 2400 * 2 + decoded = np.frombuffer(b"".join(chunks), dtype=np.int16) + assert np.all(np.diff(decoded.astype(np.int32)) >= 0) + def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture): """stream_format=sse should request delta-style multimodal outputs.""" qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 39c9e66d9..1d34b6e0b 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -56,6 +56,7 @@ def _build_test_app( include_sample_rate=False, tts_params=None, collect=None, + cumulative_audio=False, ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -386,9 +387,11 @@ class TestStreamingSpeechWebSocket: include_sample_rate=False, tts_params=None, collect=None, + cumulative_audio=False, ): assert request_start_s is not None assert request_arrival_ts is not None + assert cumulative_audio is True for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -480,6 +483,7 @@ class TestStreamingSpeechWebSocket: include_sample_rate=False, tts_params=None, collect=None, + cumulative_audio=False, ): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 7dc7a6ab3..5df9dfa7e 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -955,6 +955,17 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], """ assert response.success, "The request failed." + if request_config.get("word_timestamps"): + timestamps = response.word_timestamps + assert isinstance(timestamps, list) and timestamps, "Expected nonempty X-Word-Timestamps" + previous_start = 0 + for timestamp in timestamps: + assert isinstance(timestamp["word"], str) and timestamp["word"] + start, end = timestamp["start_ms"], timestamp["end_ms"] + assert isinstance(start, int) and isinstance(end, int) + assert previous_start <= start <= end + previous_start = start + # Optional floor on decoded audio size (models with very short clips may use a lower value). min_audio = request_config.get("min_audio_bytes") if min_audio is not None: diff --git a/tests/helpers/client.py b/tests/helpers/client.py index ed5158b92..e1042321a 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -102,6 +102,7 @@ class OmniResponse: audio_content: str | None = None audio_format: str | None = None audio_bytes: bytes | None = None + word_timestamps: list[dict[str, str | int]] | None = None #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the #: OpenAI client call through response parsing and local post-process (e.g. audio decode). e2e_latency: float | None = None @@ -1419,6 +1420,9 @@ class OnlineOmniClient: result.success = True result.audio_format = getattr(response, "response", None) if result.audio_format is not None: + timestamps = result.audio_format.headers.get("X-Word-Timestamps") + if timestamps is not None: + result.word_timestamps = json.loads(timestamps) result.audio_format = result.audio_format.headers.get("content-type", "") except Exception as e: @@ -1472,6 +1476,7 @@ class OnlineOmniClient: "extra_params", "stream_format", "x_vector_only_mode", + "word_timestamps", ): if key in request_config: extra_body[key] = request_config[key] diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index e9b1a771c..55956a6f2 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -173,3 +173,26 @@ def test_bounded_tail_allows_ratio_bound_for_long_answers(): # Twenty expected words allow up to max(2, ceil(0.2 * 20)) = 4 tail words. assert assertions._transcript_has_bounded_tail(four_word_tail, expected) assert not assertions._transcript_has_bounded_tail(five_word_tail, expected) + + +def test_speech_timestamp_header_is_required_and_decoded(): + import time + + import httpx + from openai._legacy_response import HttpxBinaryResponseContent + + from tests.helpers.client import OnlineOmniClient + + client = object.__new__(OnlineOmniClient) + request = {"word_timestamps": True} + for headers in ({}, {"X-Word-Timestamps": '[{"word":"Hello","start_ms":0,"end_ms":300}]'}): + response = client._process_non_stream_audio_speech_response( + HttpxBinaryResponseContent(httpx.Response(200, content=b"audio", headers=headers)), + wall_start=time.perf_counter(), + ) + if headers: + assert_audio_speech_response(response, request) + assert response.word_timestamps == [{"word": "Hello", "start_ms": 0, "end_ms": 300}] + else: + with pytest.raises(AssertionError, match="X-Word-Timestamps"): + assert_audio_speech_response(response, request) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d86a3216f..56b40b836 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -50,6 +50,8 @@ from vllm_omni.config.stage_config import ( merge_sampling_constraints, normalize_pipeline_cli_overrides, reconcile_diffusion_attention_overrides, + resolve_stage_async_chunk, + validate_stage_async_chunk_edges, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -1520,7 +1522,7 @@ def _build_common_stage_config_kwargs( *, model: str | None, ) -> tuple[dict[str, Any], str | None, str | None]: - input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk)) + input_proc, next_stage_proc = _select_processor_funcs(topology, resolve_stage_async_chunk(deploy, stage_deploy)) quantization_config = _build_quantization_config(deploy, engine.quantization) parallel_config = _build_parallel_config(deploy, engine.parallel, parallel_config_cls) @@ -1835,7 +1837,7 @@ def _build_connector_config( output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None return cast(Any, OmniStageConnectorConfig)( - async_chunk=bool(deploy.async_chunk), + async_chunk=resolve_stage_async_chunk(deploy, stage_deploy), omni_kv_config=_copy_value(engine.get("omni_kv_config")), output_connectors=_copy_value(output_connectors) if output_connectors else None, input_connectors=_copy_value(input_connectors) if input_connectors else None, @@ -1966,6 +1968,7 @@ class VllmOmniConfig: if len(pipeline_cfg.stages) <= 1: deploy.async_chunk = False _validate_async_chunk_support(pipeline_cfg, deploy) + validate_stage_async_chunk_edges(pipeline_cfg, deploy) deploy_by_id = {stage.stage_id: stage for stage in deploy.stages} model = cli_overrides.get("model") diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index c96f514e8..6af65b8ed 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -390,6 +390,9 @@ class StageDeployConfig: num_replicas: int = 1 env: dict[str, Any] | None = None + # False opts this stage out of pipeline-wide async chunking. + async_chunk: bool | None = None + # Inter-stage connector wiring and request defaults. output_connectors: dict[str, str] | None = None input_connectors: dict[str, str] | None = None @@ -587,6 +590,7 @@ class DeployConfig: _STAGE_RESERVED_KEYS = frozenset( { + "async_chunk", "stage_id", "devices", "num_replicas", @@ -645,6 +649,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: if name in flat_args: kwargs[name] = flat_args.pop(name) + kwargs["async_chunk"] = stage_data.get("async_chunk") kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") @@ -888,6 +893,28 @@ def _resolve_execution_mode( return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None)) +def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool: + return bool(deploy.async_chunk and (stage is None or stage.async_chunk is not False)) + + +def validate_stage_async_chunk_edges(pipeline: PipelineConfig, deploy: DeployConfig) -> None: + stages = {stage.stage_id: stage for stage in pipeline.stages} + deploy_by_id = {stage.stage_id: stage for stage in deploy.stages} + for consumer in pipeline.stages: + for source_id in consumer.input_sources: + producer = stages[source_id] + if not producer.async_chunk_process_next_stage_input_func: + continue + producer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(source_id)) + consumer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(consumer.stage_id)) + if producer_async != consumer_async: + raise ValueError( + f"Pipeline {pipeline.model_type!r} has incompatible async_chunk settings on " + f"connector edge {source_id} -> {consumer.stage_id}. " + "Set the same async_chunk mode on both stages or disable pipeline-wide async_chunk." + ) + + def _select_processor_funcs( ps: StagePipelineConfig, async_chunk: bool, @@ -964,9 +991,7 @@ def _build_engine_args( continue engine_args[k] = v engine_args.update(ds.engine_extras) - # Materialize the resolved pipeline-wide async_chunk value into every - # stage so explicit False overrides do not get lost downstream. - engine_args["async_chunk"] = bool(deploy.async_chunk) + engine_args["async_chunk"] = resolve_stage_async_chunk(deploy, ds) engine_args["session_mode"] = deploy.session_mode if deploy.session_mode == "duplex": # The engine admission limit is also the authoritative capacity for @@ -1060,11 +1085,12 @@ def merge_pipeline_deploy( "Either set async_chunk=False or implement an async-chunk producer on the pipeline." ) + validate_stage_async_chunk_edges(pipeline, deploy) result: list[StageConfig] = [] for ps in pipeline.stages: ds = deploy_by_id.get(ps.stage_id) stage_type, worker_type = _resolve_execution_mode(ps.execution_type) - input_proc, next_stage_proc = _select_processor_funcs(ps, deploy.async_chunk) + input_proc, next_stage_proc = _select_processor_funcs(ps, resolve_stage_async_chunk(deploy, ds)) engine_args = _build_engine_args(ps, ds, pipeline, deploy, next_stage_proc) # Downstream stages may share a multimodal wrapper class without owning # an encoder. Do not make vLLM profile dummy multimodal inputs for them. diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index 176a8bf69..a646a3a30 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass, field from typing import Any @@ -33,7 +33,7 @@ def get_stage_connector_role(model_config: Any) -> str | None: def stage_receives_chunks(model_config: Any) -> bool: """Whether connector chunks, rather than the orchestrator, feed a stage.""" - return get_stage_connector_role(model_config) != "sender" + return bool(getattr(model_config, "async_chunk", True)) and get_stage_connector_role(model_config) != "sender" def stage_sends_async_output(model_config: Any) -> bool: diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 87ef8d5e6..85310e9dd 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -271,8 +271,9 @@ class OmniEngineBase: ) self.num_stages = len(self.stage_configs) - stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None - self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) + self.async_chunk = any( + bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs + ) self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index f54ca3e28..80f357651 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -32,6 +32,7 @@ from vllm.entrypoints.launchers.launcher import terminate_if_errored from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector +from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError @@ -1444,14 +1445,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params: dict[str, Any] | None = None, collect: dict | None = None, target_sample_rate: int | None = None, + cumulative_audio: bool = False, ): """Generate audio chunks for streaming response. Handles two audio output modes from the engine: - Cumulative mode (list): Engine returns growing list of chunks; we emit only the new tail on each iteration. - - Per-step mode (tensor): Engine returns single tensor per iteration; - we emit it directly. + - Tensor mode: Emit each tensor directly, or only its new samples when + cumulative_audio is set for requests that retain audio for alignment. Args: generator: Async generator from the engine @@ -1462,6 +1464,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Raw audio bytes for each chunk (with WAV header for first chunk if wav format) """ prev_count = 0 + emitted_samples = 0 sample_rate_val = 24000 first_chunk = True first_audio_chunk_s: float | None = None @@ -1563,6 +1566,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) if chunk_np.ndim > 1: chunk_np = chunk_np.squeeze() + if cumulative_audio and not isinstance(audio_val, list): + total_samples = chunk_np.shape[-1] + chunk_np = chunk_np[..., emitted_samples:] + emitted_samples = total_samples + if chunk_np.size == 0: + continue if resampler is not None: chunk_np = resampler.process(chunk_np) if chunk_np.size == 0: @@ -1957,6 +1966,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY output_modalities.append(TIMESTAMPS_MODALITY) + if is_streaming_request: + # The aligner consumes the terminal waveform, so retain earlier + # audio chunks instead of draining them after each DELTA output. + sampling_params_list = [ + sp.clone() if isinstance(sp, SamplingParams) else sp for sp in sampling_params_list + ] + for sp in sampling_params_list: + if isinstance(sp, SamplingParams): + sp.output_kind = RequestOutputKind.CUMULATIVE generator = self.engine_client.generate( prompt=prompt, @@ -1990,6 +2008,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params: dict[str, Any] | None = None, collect: dict | None = None, target_sample_rate: int | None = None, + cumulative_audio: bool = False, ): """Yield raw PCM byte chunks from the engine generator. @@ -2009,6 +2028,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params=tts_params, collect=collect, target_sample_rate=target_sample_rate, + cumulative_audio=cumulative_audio, ) ) as chunks: async for chunk in chunks: @@ -2024,6 +2044,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id, tts_params=tts_params, target_sample_rate=request.sample_rate, + cumulative_audio=request.word_timestamps, ) ) as chunks: async for chunk in chunks: diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index ceeb208c9..45bdd4e90 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -483,6 +483,7 @@ class OmniStreamingSpeechHandler: include_sample_rate=True, tts_params=tts_params, collect=collect, + cumulative_audio=True, ) ) as stream: async for chunk, chunk_sample_rate in stream: diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py index 3b578ee51..e6dcecb57 100644 --- a/vllm_omni/utils/forced_aligner.py +++ b/vllm_omni/utils/forced_aligner.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Forced-aligner config + word-timestamp decoding for TTS. The aligner runs as a pooling stage appended to the pipeline by @@ -153,6 +153,7 @@ def inject_forced_aligner_stage( # Pooling is one-shot -> sync AR scheduler (the async one assumes # token-by-token generation). async_scheduling=False, + async_chunk=False, ) deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds]) logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model) From c8b3a17bd3be56ebef509987498907bcbfc0b660 Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Wed, 16 Sep 2026 17:29:57 -0700 Subject: [PATCH 15/99] [Frontend] Allow explicit turn-based serving (#7675) Signed-off-by: Linze-Shi --- docs/configuration/stage_configs.md | 2 +- docs/design/fullduplex.md | 19 ++- docs/serving/full_duplex_api.md | 30 +++- docs/serving/realtime_duplex_api.md | 47 +++---- tests/e2e/online_serving/test_minicpmo_4_5.py | 23 +++- .../openai_api/test_duplex_api_server.py | 130 ++++++++++++++---- vllm_omni/config/stage_config.py | 6 +- vllm_omni/deploy/minicpmo_4_5.yaml | 3 +- vllm_omni/deploy/minicpmo_4_5_turn.yaml | 3 + vllm_omni/entrypoints/openai/api_server.py | 38 +++-- 10 files changed, 205 insertions(+), 96 deletions(-) create mode 100644 vllm_omni/deploy/minicpmo_4_5_turn.yaml diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index aa81d42c6..199241cf6 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -53,7 +53,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze | ------- | ------ | ---------- | --------- | ------------- | | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. | | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. | -| `session_mode` | str | optional | `"turn"` | Session behavior. Duplex models (pipelines with a `duplex_plugin`, e.g. MiniCPM-o 4.5) set `"duplex"`; `vllm-omni serve` then serves it over `/v1/realtime?duplex=1` and `/v1/chat/completions`, both on duplex sessions. | +| `session_mode` | str | optional | `"turn"` | Session behavior. For pipelines with a `duplex_plugin`, explicitly select `"duplex"` for the duplex engine or `"turn"` for the ordinary online serving stack. The shipped MiniCPM-o 4.5 default remains `"duplex"`. See [Full Duplex](../serving/full_duplex_api.md#enable-full-duplex). | | `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. | | `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. | | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). | diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 4646b85cd..41a050a47 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -108,18 +108,17 @@ plugin and the session runtime config to `DuplexOrchestrator` directly. `AsyncOmniEngine`, `OrchestratorBase` and `Orchestrator` carry only template seams; `tests/engine/test_duplex_import_boundary.py` checks that importing the turn-based stack loads no duplex module. -4. **Every surface a duplex server exposes is backed by a session.** +4. **Deployment selects the serving stack.** `vllm-omni serve` constructs `DuplexOmni` when the pipeline declares - `duplex_plugin`; the server exposes `/v1/realtime?duplex=1` (alias + `duplex_plugin` and the deploy configuration sets `session_mode: duplex`. + An explicit `session_mode: turn` selects `AsyncOmni` instead, without + changing the pipeline's duplex capability. In duplex mode the server exposes `/v1/realtime?duplex=1` (alias `/v1/duplex`), `POST /v1/chat/completions`, `/v1/models` and `/health`, and - every other turn-based route reports "not available". The chat route is - served by `DuplexChatCompletionsAdapter`, a Realtime client that runs one - short-lived session per request: speech becomes a committed turn, text is - seeded as the session's opening turn because a model-native model takes a - turn only when it hears speech. Whether a model can be reached with text is - `DuplexCapabilities.supports_chat_completions`; one that cannot is refused - at once rather than left to idle out. Turn-based use of the same model stays - available offline through `Omni` / `AsyncOmni`. + every other turn-based route reports "not available". The chat route uses + the ordinary chat service on the duplex engine, without a session, when + the plugin declares `DuplexCapabilities.supports_chat_completions`. + Turn-based deployments use the ordinary API initialization and retain + model capability checks and endpoint restrictions. 5. **Serving is transport only.** The websocket handler does socket I/O, wire-envelope validation, command translation, event rendering and the attachment/resume/replay bookkeeping; it holds no session state. diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md index 92511e84f..9b3eceb58 100644 --- a/docs/serving/full_duplex_api.md +++ b/docs/serving/full_duplex_api.md @@ -22,14 +22,14 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md). ## Enable Full Duplex A model is served full duplex when its registered pipeline declares a -`duplex_plugin` (the model's `DuplexModelPlugin`). That declaration alone -decides it: `vllm-omni serve` constructs `DuplexOmni` instead of `AsyncOmni`, -and every surface the server exposes is backed by a duplex session. It serves +`duplex_plugin` (the model's `DuplexModelPlugin`) and its deploy configuration +sets `session_mode: duplex`. `vllm serve --omni` constructs `DuplexOmni` +instead of `AsyncOmni`. It serves `/v1/realtime?duplex=1` (and its alias `/v1/duplex`), `POST /v1/chat/completions`, `/v1/models` and `/health`; every other turn-based HTTP route (speech, batch, embeddings, video, ...) answers "not -available". Turn-based use of such a model stays available offline through the -Python API (`Omni` / `AsyncOmni`). +available". Set `session_mode: turn` to select the ordinary online serving +stack instead. The model's supported tasks and endpoint restrictions still apply. `/v1/chat/completions` is the ordinary chat service running on the duplex engine: a request is a turn-based generation on the same stages, served @@ -45,8 +45,24 @@ The deploy configuration of such a model must agree: session_mode: duplex ``` -A duplex model started with a deploy configuration that does not set it fails -at startup rather than falling back to turn-based serving. +A duplex-capable model must explicitly set `session_mode` to `duplex` or `turn` +in its deploy configuration, possibly through `base_config` inheritance. +Missing or invalid values fail at startup. The default MiniCPM-o deployment +continues to use duplex mode. + +To run MiniCPM-o 4.5 with the turn-based engine: + +```bash +vllm serve openbmb/MiniCPM-o-4_5 --omni \ + --deploy-config vllm_omni/deploy/minicpmo_4_5_turn.yaml \ + --trust-remote-code \ + --port 8091 +``` + +This profile inherits the default model and stage settings and overrides only +`session_mode`. It supports ordinary HTTP requests without creating a duplex +session handler. Mode selection happens at startup; a WebSocket query parameter +does not switch engines. The Python `Omni` / `AsyncOmni` APIs are unchanged. !!! warning diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index cc10d99c9..c18b416c7 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -32,33 +32,19 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \ mounts `ws://:8099/v1/realtime?duplex=1` (this page; `ws://:8099/v1/duplex` is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and `/health`, and nothing else. The remaining turn-based HTTP routes (speech, -batch, embeddings, video, ...) are not served by a duplex model. +batch, embeddings, video, ...) are not served in duplex mode. -`/v1/chat/completions` is not a turn-based path on a duplex server: each -request opens a short-lived duplex session and reads the answer off it. Audio -content becomes an ordinary committed turn; text is seeded as the session's -opening turn, because a model-native model takes a turn only when it hears -speech. A model that cannot be seeded that way declares so, and a text request -to it is refused with 400. Two consequences are worth sizing for: +`/v1/chat/completions` uses the ordinary chat service on the same engine +when the plugin declares `supports_chat_completions`. HTTP chat requests do +not create duplex sessions or consume session admission slots. Their options +follow the [Chat Completions API](chat_completions_api.md) and the model's +capabilities and endpoint restrictions. -- **Every request holds an admission slot** for its lifetime, so - `duplex_session.max_sessions` caps HTTP concurrency as well as websocket - sessions. Beyond it a request is refused with HTTP 503. -- **Latency follows the model's own clock.** A model-native session generates - per audio unit, so even a short answer costs what that many units cost -- - noticeably more than the same model would take turn-based. - -`n > 1`, `logprobs`, `tools`, and image or video content parts have no -representation in one duplex turn and are refused with HTTP 400 -- a Realtime -conversation item carries text and audio only. `temperature`, -`max_completion_tokens`, `modalities` and -`chat_template_kwargs.use_tts_template` apply; the other -`chat_template_kwargs` keys are logged as ignored. `extra_body` is passed to -the session unchanged, exactly as a websocket client's session `extra_body` -is, which is how a request asking for `modalities: ["text", "audio"]` supplies -the `ref_audio` MiniCPM-o requires for audio output. A model that should not -serve the route lists it in the deploy config's `endpoint_restrictions`, -exactly as a turn-based model does. +For a turn-based online deployment, use `minicpmo_4_5_turn.yaml` or explicitly +set `session_mode: turn` in a deploy configuration derived from the model's +profile. This selects `AsyncOmni` and the ordinary API initialization instead +of the duplex handler. See [Full Duplex](full_duplex_api.md#enable-full-duplex) +for the command and configuration requirements. ### Run the Example Client @@ -1345,8 +1331,9 @@ them out into typed events before they reach a client. - Session capacity is bounded by `duplex_session.max_sessions` in the deploy configuration; admission beyond it fails with `resource_exhausted` (`rate_limit_error`, retryable). -- A duplex deployment serves the duplex session surfaces only: the websocket - route and `/v1/chat/completions`, the latter on a session per request (see - [Start the Server](#start-the-server)). The other turn-based HTTP routes - report "not available" on a MiniCPM-o 4.5 server; turn-based use of the - model stays available offline through `Omni` / `AsyncOmni`. +- A duplex deployment serves the websocket route and, when supported by the + model, ordinary `/v1/chat/completions` requests on the same engine without a + session per request (see [Start the Server](#start-the-server)). Other + turn-based HTTP routes report "not available" in duplex mode. To use the + ordinary online serving stack, select `session_mode: turn` as described in + [Full Duplex](full_duplex_api.md#enable-full-duplex). diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index 5417d2b16..a1d045b72 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5.py @@ -83,6 +83,25 @@ test_params = [ ), ] +turn_test_params = [ + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=modify_stage_config(_CI_DEPLOY, updates={"session_mode": "turn"}), + use_stage_cli=False, + server_args=[ + "--trust-remote-code", + "--async-chunk", + "--chat-template", + str(_NATIVE_CHAT_TEMPLATE), + "--chat-template-content-format", + "openai", + ], + ), + id="turn_async_chunk", + ), +] + def get_system_prompt(): return { @@ -138,7 +157,7 @@ def _read_prompt_selection_log(offset: int) -> str: @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("omni_server", test_params + turn_test_params, indirect=True) def test_text_to_text_001(omni_server, openai_client) -> None: """ Test text-only input generating text output via OpenAI API. @@ -163,7 +182,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.full_model @pytest.mark.omni @hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("omni_server", test_params + turn_test_params, indirect=True) def test_text_to_audio_001(omni_server, openai_client) -> None: """ Test text-only input generating text + audio output via OpenAI API. diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py index b528cbc2a..9cf58b114 100644 --- a/tests/entrypoints/openai_api/test_duplex_api_server.py +++ b/tests/entrypoints/openai_api/test_duplex_api_server.py @@ -1,28 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Guards for the duplex-only API server (unified full-duplex framework). - -A model whose pipeline declares ``duplex_plugin`` is served through -``DuplexOmni`` and nothing else: sessions run over ``/v1/realtime?duplex=1`` -(alias ``/v1/duplex``) and every turn-based route reports "not available". -These tests lock that shape, which ``test_api_server_guards.py`` covers for the -diffusion and multi-stage servers: - -* the pipeline probe that decides duplex vs turn-based at startup, -* the ``app.state`` snapshot of a duplex server (what is wired, what stays - ``None``), -* both duplex websocket routes reaching the session handler, and - ``/v1/realtime`` without the flag still reporting the turn-based Realtime API - as unavailable, -* the startup-warmup gate that holds real clients while the warmup session runs. - -CPU-only: no engine is started; ``DuplexOmni`` is subclassed with just the -accessors the app-state wiring reads. -""" +"""CPU tests for serving-mode selection, duplex API wiring and warmup.""" from __future__ import annotations import asyncio +from argparse import Namespace +from dataclasses import replace from types import SimpleNamespace import pytest @@ -31,12 +15,14 @@ from fastapi.testclient import TestClient from starlette.datastructures import State from starlette.websockets import WebSocketDisconnect +from vllm_omni.config import stage_config from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig, StagePipelineConfig from vllm_omni.engine.duplex.config import DuplexCapabilities from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex_omni import DuplexOmni from vllm_omni.entrypoints.openai import api_server +from vllm_omni.utils.tracking_parser import TrackingNamespace pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -202,9 +188,10 @@ def _minimal_args(**overrides) -> SimpleNamespace: (None, False), ], ) -def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_config, expected: bool) -> None: - """``duplex_plugin`` alone decides whether ``serve`` builds ``DuplexOmni``.""" +def test_duplex_serving_requires_a_pipeline_plugin(monkeypatch, tmp_path, pipeline_config, expected: bool) -> None: seen: dict[str, object] = {} + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text("session_mode: duplex\n") def fake_get_pipeline_config(*, model, trust_remote_code, deploy_config_path=None, **_kwargs): seen.update(model=model, trust_remote_code=trust_remote_code, deploy_config_path=deploy_config_path) @@ -212,12 +199,12 @@ def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_co monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) - kwargs = {"trust_remote_code": True, "deploy_config": "deploy.yaml"} - assert api_server._is_duplex_model("demo-duplex-model", kwargs) is expected + kwargs = {"trust_remote_code": True, "deploy_config": str(deploy_path)} + assert api_server._should_serve_duplex("demo-duplex-model", kwargs) is expected assert seen == { "model": "demo-duplex-model", "trust_remote_code": True, - "deploy_config_path": "deploy.yaml", + "deploy_config_path": str(deploy_path), } @@ -230,7 +217,98 @@ def test_duplex_model_probe_propagates_resolution_errors(monkeypatch) -> None: monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) with pytest.raises(ValueError, match="trust_remote_code"): - api_server._is_duplex_model("demo-duplex-model", {}) + api_server._should_serve_duplex("demo-duplex-model", {}) + + +@pytest.fixture +def duplex_pipeline(): + return PipelineConfig( + model_type="demo", + stages=(StagePipelineConfig(stage_id=0, model_stage="demo", final_output=True),), + duplex_plugin="pkg.mod.Plugin", + ) + + +@pytest.mark.parametrize("session_mode", ["turn", "duplex"]) +@pytest.mark.parametrize("config_source", ["explicit", "default", "basename", "inherited"]) +@pytest.mark.asyncio +async def test_serving_mode_selects_engine_and_preserves_stage_config( + monkeypatch, tmp_path, duplex_pipeline, session_mode, config_source +): + base_path = tmp_path / "base.yaml" + base_path.write_text("session_mode: duplex\nstages:\n - stage_id: 0\n max_num_seqs: 4\n") + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text(f"base_config: base.yaml\nsession_mode: {session_mode}\n") + if config_source == "inherited": + deploy_path = tmp_path / "overlay.yaml" + deploy_path.write_text("base_config: deploy.yaml\n") + pipeline = replace(duplex_pipeline, default_deploy_config_name=deploy_path.name) + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: pipeline) + monkeypatch.setattr(stage_config, "_DEPLOY_DIR", tmp_path) + kwargs = {} + if config_source != "default": + kwargs["deploy_config"] = deploy_path.name if config_source == "basename" else str(deploy_path) + + created = [] + + class TurnEngine: + def __init__(self, model, **engine_kwargs): + self.model = model + self.kwargs = engine_kwargs + self.closed = False + created.append(self) + + def shutdown(self): + self.closed = True + + class DuplexEngine(TurnEngine): + pass + + monkeypatch.setattr(api_server, "AsyncOmni", TurnEngine) + monkeypatch.setattr(api_server, "DuplexOmni", DuplexEngine) + args = TrackingNamespace( + Namespace(model="demo-model", disable_log_stats=True, **kwargs), frozenset({"model", *kwargs}) + ) + async with api_server.build_async_omni_from_stage_config(args) as engine: + assert type(engine) is (DuplexEngine if session_mode == "duplex" else TurnEngine) + assert engine.model == "demo-model" + assert engine.kwargs == {**kwargs, "log_stats": False} + assert not engine.closed + assert len(created) == 1 + assert engine.closed + deploy = stage_config.load_deploy_config(deploy_path) + assert deploy.session_mode == session_mode + assert deploy.stages[0].max_num_seqs == 4 + + +@pytest.mark.parametrize("contents", ["{}\n", "session_mode: null\n", "session_mode: typo\n"]) +def test_duplex_serving_rejects_missing_or_invalid_mode(monkeypatch, tmp_path, duplex_pipeline, contents): + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text(contents) + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline) + with pytest.raises(ValueError, match="session_mode: turn or duplex"): + api_server._should_serve_duplex("demo-model", {"deploy_config": str(deploy_path)}) + + +def test_duplex_serving_requires_a_deploy_config(monkeypatch, duplex_pipeline): + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline) + with pytest.raises(ValueError, match="requires a deploy config"): + api_server._should_serve_duplex("demo-model", {}) + + +def test_duplex_serving_propagates_deploy_load_errors(monkeypatch, tmp_path, duplex_pipeline): + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline) + with pytest.raises(FileNotFoundError): + api_server._should_serve_duplex("demo-model", {"deploy_config": str(tmp_path / "missing.yaml")}) + + +@pytest.mark.parametrize("deploy_name,expected", [("minicpmo_4_5.yaml", True), ("minicpmo_4_5_turn.yaml", False)]) +def test_minicpmo_serving_profiles(monkeypatch, deploy_name, expected): + from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: MINICPMO_4_5_PIPELINE) + assert api_server._should_serve_duplex("openbmb/MiniCPM-o-4_5", {"deploy_config": deploy_name}) is expected + assert api_server._should_serve_duplex("openbmb/MiniCPM-o-4_5", {}) is True # --------------------------------------------------------------------------- # diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 6af65b8ed..24453ae88 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -309,10 +309,8 @@ class PipelineConfig: diffusers_class_name: str | None = None diffusers_class_aliases: tuple[str, ...] = () endpoint_restrictions: tuple[EndpointRestriction, ...] = () - # Dotted path of the model's ``DuplexModelPlugin``. A pipeline is a duplex - # model iff this is set: ``vllm-omni serve`` then always runs it through - # ``DuplexOmni`` (every served surface runs on a duplex session) and the engine hosts one - # ``DuplexOrchestrator`` with the plugin loaded. + # Dotted path of the model's ``DuplexModelPlugin``. Online serving uses + # DuplexOmni only when the deploy configuration selects session_mode: duplex. duplex_plugin: str | None = None # Legacy duplex wiring of the models that are not ported to the plugin # framework yet (PersonaPlex, Nemotron VoiceChat). Nothing reads them: a diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 5c26fc2c6..a04cff400 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,8 +1,7 @@ # MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. pipeline: minicpmo_4_5 async_chunk: true -# session_mode: duplex makes vllm-omni serve run this model duplex-only -# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. +# Use minicpmo_4_5_turn.yaml to select the ordinary turn-based serving engine. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_turn.yaml b/vllm_omni/deploy/minicpmo_4_5_turn.yaml new file mode 100644 index 000000000..8dada1257 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_turn.yaml @@ -0,0 +1,3 @@ +# MiniCPM-o 4.5 using the turn-based engine and ordinary HTTP serving. +base_config: minicpmo_4_5.yaml +session_mode: turn diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 376a9ccd2..20706396c 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -19,6 +19,7 @@ from argparse import Namespace from collections.abc import AsyncIterator from contextlib import asynccontextmanager from http import HTTPStatus +from pathlib import Path from typing import Annotated, Any, Literal import vllm.envs as envs @@ -440,21 +441,34 @@ async def build_async_omni( yield async_omni -def _is_duplex_model(model: str, kwargs: dict[str, Any]) -> bool: - """Whether ``model``'s pipeline declares a ``duplex_plugin`` (served by ``DuplexOmni``). - - Resolution errors propagate: a duplex model whose pipeline or deploy - config cannot be resolved must fail startup rather than silently start a - turn-based server. - """ +def _should_serve_duplex(model: str, kwargs: dict[str, Any]) -> bool: + """Select the serving engine without changing the model's duplex capability.""" from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.config.stage_config import _DEPLOY_DIR, resolve_deploy_yaml pipeline_config = StageConfigFactory.get_pipeline_config( model=model, trust_remote_code=bool(kwargs.get("trust_remote_code")), deploy_config_path=kwargs.get("deploy_config"), ) - return bool(pipeline_config is not None and getattr(pipeline_config, "duplex_plugin", None)) + if pipeline_config is None or not getattr(pipeline_config, "duplex_plugin", None): + return False + + deploy_path = kwargs.get("deploy_config") + if deploy_path is None: + if pipeline_config.default_deploy_config_name is None: + raise ValueError("A duplex-capable model requires a deploy config with session_mode: turn or duplex") + deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name + else: + deploy_path = Path(deploy_path) + if not deploy_path.exists() and deploy_path.parent == Path("."): + deploy_path = _DEPLOY_DIR / deploy_path + + # Resolve base_config too, so the API and stage workers use the same mode. + session_mode = resolve_deploy_yaml(deploy_path).get("session_mode") + if session_mode not in ("turn", "duplex"): + raise ValueError("A duplex-capable model requires session_mode: turn or duplex in its deploy config") + return session_mode == "duplex" @asynccontextmanager @@ -518,9 +532,7 @@ async def build_async_omni_from_stage_config( kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) - if _is_duplex_model(model, kwargs): - # A duplex model is always served in duplex mode: sessions over - # /v1/realtime?duplex=1, no turn-based HTTP routes. + if _should_serve_duplex(model, kwargs): async_omni = DuplexOmni(model=model, **kwargs) else: async_omni = AsyncOmni(model=model, **kwargs) @@ -712,9 +724,7 @@ async def omni_init_app_state( state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None model_name = served_model_names[0] if served_model_names else args.model - # Duplex mode: a duplex model is served through DuplexOmni only. Sessions - # run over /v1/realtime?duplex=1; every turn-based HTTP route reports - # "not available". + # Initialize the API surface for the selected engine, not just the model. if isinstance(engine_client, DuplexOmni): await _init_duplex_app_state(engine_client, state, args, base_model_paths, vllm_config, request_logger) return From d81777f297811ceb68bd991b6defd9f46fc9162e Mon Sep 17 00:00:00 2001 From: Bojiang Li <327132355+bojiang-li@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:30:47 -0700 Subject: [PATCH 16/99] [Doc] Scope Spark FP8 guidance to the validated revision (#7536) Signed-off-by: bojiang-li <327132355+bojiang-li@users.noreply.github.com> --- recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md index efd88c068..f4304cd4a 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md +++ b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md @@ -16,6 +16,12 @@ Validated on: 50 steps, text-encoder TP2, and DiT USP2 - Common sampling settings: `flow_shift=12`, `seed=1101` +The commands and measurements below describe this pinned revision. Newer +versions have changed the scope of `--quantization fp8`; see +[Online FP8 quantization](MiniMax-H3.md#online-fp8-quantization) for the current +component selection and offload support. Re-qualify memory, output, and latency +on GB10 before comparing results across revisions. + ## Capacity requirements | Resource | DGX Spark (GB10) | @@ -40,8 +46,10 @@ math relative to the RTX and datacenter recipes: - **Do not use `--enable-cpu-offload`.** Moving weights from "VRAM" to "host RAM" frees nothing here. - **`--quantization fp8` is mandatory.** A BF16 partition is 135 GiB and does not - fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the - Qwen3-VL text encoder and both VAEs stay BF16. + fit in 121 GiB. At the validated revision, online FP8 quantizes the DiT only + (62 GiB to ~31 GiB); the Qwen3-VL text encoder and both VAEs stay BF16. + Current versions also quantize eligible text-decoder linears by default, so + the same flag does not imply the same quantization scope across revisions. With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly 23 GiB of headroom for the OS, page cache, and any other process on the box. @@ -295,6 +303,8 @@ shorter than a 50-step Ref2VA run. - Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note under **Validated on**: FP8 weight loading and image-only Ref2VA are both broken on `release/v0.26.0`. The suggested version is `v0.26.1`. -- Online FP8 is incompatible with layerwise offload — the offload path produces a - weight stride the Cutlass FP8 kernel rejects. This is not a practical - restriction here since offload is unusable on GB10 anyway. +- At the validated revision, online FP8 with layerwise offload produced a weight + stride the Cutlass FP8 kernel rejected. Current H3 supports online FP8 with + layerwise offload and both DLO transfer paths; see the linked FP8 section + above. This does not change the GB10 capacity guidance: offload uses the same + unified-memory pool and is not enabled in this recipe. From 72cc012cfdf2ae531338cb9e221804f3bc31e467 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:52:23 +0800 Subject: [PATCH 17/99] [CI][bugfix] Fix MiniMax H3 FP8 quality test remote-code loading (#7617) --- .../models/minimax_h3/test_minimax_h3_quantization_quality.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index ba0e0ddf5..87b04613a 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -189,6 +189,8 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig): # The fused BF16 baseline only fits on H100-80GB with encoder TP. "stage_0_tensor_parallel_size": 2, "vae_use_tiling": True, + # MiniMax H3's VAE is implemented by code shipped with the checkpoint. + "trust_remote_code": True, } omni_bl = Omni(**common_kwargs) From 4fab34c8f46f1f431571cf5c699f55753b4366ea Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Wed, 16 Sep 2026 19:45:31 -0700 Subject: [PATCH 18/99] [Bugfix] Update ComfyUI serving-mode mock after helper rename (#7694) Signed-off-by: Linze-Shi --- tests/e2e/features/comfyui/test_comfyui_integration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index 3c6f0492a..480381f0a 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -534,7 +534,7 @@ def mock_async_omni( # 1. The API layer uses its stage_list and stage_configs attributes # 2. Its __init__ method has slow side effects (model & config loading). # These cases use the non-duplex mock engine without loading model configs. - mocker.patch("vllm_omni.entrypoints.openai.api_server._is_duplex_model", return_value=False) + mocker.patch("vllm_omni.entrypoints.openai.api_server._should_serve_duplex", return_value=False) mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni") monkeypatch.setattr( "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat", From becfefd4402da180d8af5121cef474fee40883af Mon Sep 17 00:00:00 2001 From: HAAZZZEEEE <49724853+HAAZZZEEEE@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:11:40 +0800 Subject: [PATCH 19/99] [Hardware][Ascend] Add Wan2.2 MXFP4 UOS, Smooth and W4A8 fallback (#7210) Signed-off-by: HAAZZZEEEE --- docs/user_guide/quantization/mxfp4.md | 414 +++++++++---- .../text_to_video/text_to_video.py | 11 +- tests/config/test_omni_config.py | 30 +- .../wan2_2/test_create_transformer_quant.py | 213 ++++++- .../wan2_2/test_wan22_pipeline_diffuse.py | 98 +++- .../quantization/test_mxfp4_checkpoint.py | 294 ++++++++++ .../quantization/test_mxfp4_config.py | 555 +++++++++++++++++- .../quantization/test_wan_mxfp4_loading.py | 237 ++++++++ .../test_diffusion_config_propagation.py | 6 + .../test_text_to_video_prompts.py | 49 +- .../npu/quant/test_mxfp4_step_fallback.py | 168 ++++++ vllm_omni/config/omni_config.py | 4 + vllm_omni/diffusion/data.py | 4 + .../model_loader/diffusers_loader.py | 37 +- .../models/wan2_2/pipeline_wan2_2.py | 51 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 26 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 11 +- .../models/wan2_2/wan2_2_transformer.py | 106 +++- vllm_omni/quantization/factory.py | 77 ++- vllm_omni/quantization/mxfp4_config.py | 228 ++++++- 20 files changed, 2435 insertions(+), 184 deletions(-) create mode 100644 tests/diffusion/quantization/test_mxfp4_checkpoint.py create mode 100644 tests/diffusion/quantization/test_wan_mxfp4_loading.py create mode 100644 tests/platforms/npu/quant/test_mxfp4_step_fallback.py diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index b18b3e7b3..00ff97c01 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -6,22 +6,25 @@ W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4 (`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale. +On Ascend, single-scale MXFP4 also supports Universal Optimal Scaling +(UOS, `Qmax=7.25`), described in the single-scale scaling section below. + vLLM-Omni provides two quantization methods with different scale structures: | Method | Scale structure | Mode | Use case | -|--------|----------------|------|----------| -| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed | -| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended | +| -------- | ---------------- | ------ | ---------- | +| `mxfp4` | Single-scale (per-32 fine only) | Online + pre-quantized single-scale weights | Online baseline, or an existing single-scale checkpoint in the layout below | +| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | W4A4 with dual-level scales; offline mode loads calibrated Smooth tensors | -!!! tip "Recommended: `mxfp4_dualscale` offline" - For production deployments, use the `mxfp4_dualscale` offline mode with a +!!! tip "Existing W4A4-only DualScale deployments" + For W4A4-only deployments, the existing `mxfp4_dualscale` offline mode uses a pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load - calibrated `mul_scale` tensors from disk, providing measurably better accuracy - than any online method. The one-time preprocessing cost amortises across all - subsequent inference runs. + calibrated `mul_scale` tensors from disk. Quality depends on the model, + calibration recipe and BF16 fallback policy; evaluate the checkpoint against + your quality target before deployment. - Use `mxfp4` online only for quick experimentation where preprocessing time - is not acceptable and accuracy loss is tolerable. + W4A8 layer/step fallback requires single-scale `mxfp4`, online or offline. + DualScale supports W4A4 and BF16 routing only. !!! warning "Online single-scale ≠ Offline dual-scale" `mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`: @@ -34,61 +37,244 @@ vLLM-Omni provides two quantization methods with different scale structures: ## Hardware Support -| Device | Support | -|--------|---------| -| NVIDIA Blackwell GPU (SM 100+) | ⭕ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | -| NVIDIA Ampere GPU (SM 80+) | ⭕ | -| AMD ROCm (gfx950 / MI355X) | ✅ | -| Intel XPU | ⭕ | -| Ascend NPU (Atlas 950 A5) | ✅ | +The following matrix describes these vLLM-Omni implementations, rather than +hardware support for the MXFP4 format in general. -Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. +| Device | Single-scale online OCP | Single-scale offline | UOS / W4A8 fallback | DualScale | +| -------- | -------- | -------- | -------- | -------- | +| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ | ✅ | +| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ | ❌ | +| Other AMD architectures | ❌ | ❌ | ❌ | ❌ | +| NVIDIA GPU | ❌ | ❌ | ❌ | ❌ | +| Intel XPU | ❌ | ❌ | ❌ | ❌ | + +Legend: `✅` supported by the implementation, `❌` unsupported. +ROCm uses the existing AITER online path and requires `mxfp4_scale_alg=0`. ## Model Type Support ### Diffusion Model (Wan2.2) | Model | Online | Offline | Notes | -|-------|--------|---------|-------| -| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config | -| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B | -| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss | +| ------- | -------- | --------- | ------- | +| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Cascade experts resolve checkpoint metadata independently; explicit global runtime policy applies to both | +| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Per-expert configuration uses the same interface as T2V; validate I2V checkpoints separately | +| Wan2.2-TI2V-5B | Requires validation | Requires validation | This guide does not establish checkpoint compatibility or generation quality for TI2V-5B | -The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about -quantization quality, not model compatibility — both work on cascade (A14B) and -single-transformer models alike, the same as `mxfp8` online: +Both methods can be integrated with single-transformer and cascade models. +This alone does not establish support or quality for every model variant: -- `mxfp4`: single-scale, lower overhead, simpler compute, online only -- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline +- `mxfp4`: single-scale; online and pre-quantized weights, with optional W4A8 fallback on NPU +- `mxfp4_dualscale`: dual-scale; online and offline, with BF16 fallback -**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced -by the merge script); there is no offline `mxfp4` single-scale format. +The existing `merge_mxfp4_dualscale_checkpoint.py` produces **offline +`mxfp4_dualscale`** checkpoints. Single-scale weights must satisfy the separate +layout described below; changing a DualScale checkpoint's method name to +`mxfp4` does not convert its weights. Prepare single-scale checkpoints with +external tooling. Real exported checkpoint compatibility and model quality +require validation of that checkpoint; the synthetic tests do not establish +UOS export or video acceptance. !!! note "Per-layer BF16 fallback in offline cascade models" - The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most + An A14B offline DualScale checkpoint uses `quant_method: mxfp4_dualscale`. Most linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers retain their original BF16 weights and are listed in `ignored_layers` inside each transformer's `config.json`. The two transformers may have different `ignored_layers` sets — the pipeline reads each transformer's own `config.json` - and rebuilds the config locally when they differ, so routing is always - per-transformer-accurate. - -!!! warning "TI2V-5B not supported" - Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter - count makes it significantly more sensitive to 4-bit quantization noise, - resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B. + and resolves the storage format and BF16 layer policy independently. ## Configuration +### W4A8 fallback at selected layers and denoising steps (Ascend NPU) + +For single-scale `mxfp4`, set `w4a8_fallback_steps` to a list of zero-based +request denoising loop indices (not scheduler timestep values). +Set `w4a8_fallback_layers` to exact runtime Linear paths, for example +`["blocks.10.attn1.to_qkv", "blocks.12.ffn.net_2"]`. A selected layer uses W4A8 +at every step; at a selected step, all other quantized layers also use W4A8. +The rules are combined with **OR**. W4A8 uses MXFP8 activations with MXFP4 +weights. Omitting both lists or passing `[]` keeps the existing W4A4 behavior. Existing +`ignored_layers` take priority: **BF16 > (layer OR step) W4A8 > W4A4**. +DualScale's existing leading BF16 block rule is unchanged; it does not enable +W4A8. + +```python +from vllm_omni import Omni + +omni = Omni( + model="", + quantization_config={ + "method": "mxfp4", + "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"], + "w4a8_fallback_steps": [0, 1, 38, 39], + }, +) +``` + +Both policies also work with offline single-scale `mxfp4`. Any nonempty +W4A8 list with `mxfp4_dualscale`, online or offline, raises an explicit error +at configuration construction. Empty lists preserve its existing W4A4 path. +Each expert's checkpoint still determines its serialized +format and BF16 layer routing. A serialized checkpoint config is the complete +storage contract: `ignored_layers`, or its `modules_to_not_convert` alias, +selects the BF16 layers, and omitting both means no explicit BF16 layers. A +method-only checkpoint value such as `"mxfp4"` or +`{"quant_method": "mxfp4"}` carries no storage policy and leaves a caller's +online `ignored_layers` unchanged. An explicitly supplied quantization config +retains its W4A8 layer and step lists when the loader rebuilds either expert's +config from disk, including empty lists that disable checkpoint policies. +Within an explicit runtime config, an omitted W4A8 list defaults to `[]`; +individual omitted fields do not inherit saved policies. Without an explicit +runtime config, auto-detection uses the saved checkpoint policies. + +Layer names are relative to each expert (`blocks.10...`, without a +`transformer.` or `transformer_2.` prefix). A global config applies the same +list to both experts; per-component `transformer`/`transformer_2` configs can +select different lists. Wan selects these expert names exactly: a `transformer` +entry does not match `transformer_2`; an omitted expert uses `default` if set. +A component mapped to `null` (or an unmatched component +without a default) stays unquantized and requires a BF16 checkpoint; loading a +quantized checkpoint for that disabled component raises an error. A top-level +`quantization_config=None` retains checkpoint auto-detection. +Names are exact, with no regex, wildcard, block expansion, +or checkpoint-name remapping. Wan validates them against its runtime Linear +modules and logs the selected paths/count and BF16 overrides per PP rank. +Use `attn1.to_qkv` for fused self-attention; `attn1.to_q`, `to_k`, or `to_v` +checkpoint names raise an error. Cross-attention `attn2.to_q/to_k/to_v` remains +separate. Layer-only fallback does not require a denoise ForwardContext. + +The text-to-video example accepts the JSON directly: + +```bash +ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 \ +python examples/offline_inference/text_to_video/text_to_video.py \ + --model /path/to/Wan2.2-T2V-A14B-Diffusers \ + --quantization-config '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}' \ + --tensor-parallel-size 1 --ulysses-degree 2 --cfg-parallel-size 2 \ + --enforce-eager --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 40 --seed 42 --output wan22-w4a8.mp4 +``` + +Use either `--quantization` or `--quantization-config` in this example. +Serving accepts the same JSON through `--diffusion-quantization-config`. + +The list uses one request's loop index across both high- and low-noise experts, +and both CFG branches see the same index. A new request starts at step zero. +Indices outside the request's step range are never selected. A pipeline must +publish `ForwardContext.denoise_step_idx`; requesting step fallback without +that context raises an error instead of silently ignoring the list. + +Single-scale W4A4 and W4A8 reuse the same prepared `weight` and `weight_scale`. +Online loading quantizes the original BF16/FP16 weight once; offline loading +packs the numeric FP4 values and reshapes the saved E8M0 scale without +requantizing the weight. Only activation precision changes during forward. +There is no DualScale decoding, second W4 representation or `w4a8_weight` / +`w4a8_weight_scale` cache. No weight quantization runs inside the denoising loop. + +The mixed FP8/FP4 GEMM uses BF16 output and a matching BF16 row bias for +compatibility with the target A5 runtime. An FP16 caller receives that result +converted back to FP16; its output shape and dtype are preserved. + +The single-scale fallback validation target is eager Wan2.2 T2V A14B on A5. +Quality scores, optimal fallback lists, real UOS exports and additional +compile/offload combinations require separate evaluation. Existing W4A4-only +DualScale still needs complete 512-channel input groups; use TP=1 with +sequence/CFG parallelism for its initial multi-card run. + +### `mxfp4` — Pre-Quantized Single-Scale Weights (Ascend NPU) + +Set the following fields in each transformer's `config.json` when a +single-scale checkpoint uses calibrated Smooth scaling: + +```json +{ + "quantization_config": { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": true, + "require_smooth_scale": true, + "ignored_layers": [] + } +} +``` + +Each quantized Linear must provide tensors under its model parameter name: + +| Tensor | Checkpoint layout | Meaning | +| -------- | ------------------- | --------- | +| `weight` | `(N, K)` numeric FP4 values | Packed to FP4 after loading | +| `weight_scale` | `(N, ceil(K / 32))`, `uint8` | Raw E8M0 exponent bytes, required | +| `mul_scale` | `(K,)`, FP32 | Smooth pre-scale, required when declared above | + +Both W4A4 and W4A8 quantize `x * mul_scale`. The Smooth tensor is converted +to the layer's activation dtype after loading. Row-parallel layers load the +corresponding input-channel shard. Pre-fused QKV weights must use one shared +Smooth tensor for the fused projection. + +For backward compatibility, `require_smooth_scale` defaults to `false`: +an absent tensor uses identity scaling, while a present tensor is still loaded. +Without a declaration, an absent Smooth tensor cannot be distinguished from a +legacy unsmoothed checkpoint. Calibrated checkpoints should declare +`require_smooth_scale: true` so missing or misnamed Smooth tensors fail loading; +incorrect tensor shapes, non-floating tensors, NaN/Inf and nonpositive values +also fail. Values that overflow or underflow to zero in the activation dtype +are rejected during weight processing. The flag requires serialized single-scale +weights and is rejected for online quantization. It does not make DualScale +scales optional or change missing-parameter policies for other backends. +Either the checkpoint or an explicit caller can require Smooth: a `true` +declaration is retained when per-expert metadata rebuilds the configuration, +including when the other source omits the flag or sets it to `false`. + +Prepare the canonical numeric FP4 checkpoint before inference using external +tooling. Runtime loading does not parse producer-specific export formats. +Retain the producer recipe to identify the weight-generation algorithm and +Smooth calibration; tensor storage alone does not establish UOS provenance. + +### `mxfp4` — Single-Scale Block Scaling (OCP / UOS) + +`mxfp4_scale_alg` selects the MXFP4 block-scaling method on NPU. +UOS is a scaling rule, not a separate activation dtype or an activation-only +algorithm. Both choices use FP4 E2M1 values with shared E8M0 scales. + +| Mode | Weight preparation | A4 activation quantization | W4A8 fallback | +| --- | --- | --- | --- | +| Online | Selected scaling method, once at loading | Selected scaling method | Same W4; OCP MXFP8 A8 | +| Offline | Load existing W4 and scales unchanged | Selected scaling method | Same W4; OCP MXFP8 A8 | + +The available settings are: + +| Setting | MXFP4 scaling method | NPU parameters | +| --- | --- | --- | +| `0` (default) | OCP MX | `scale_alg=0` | +| `2` | UOS (`Qmax=7.25`) | `scale_alg=2, dst_type_max=7.25` | + +Universal Optimal Scaling (UOS, `Qmax=7.25`) is derived in +[MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://arxiv.org/abs/2607.24377), +Section 4.1. The value `7.25` is the boundary used to select the shared +power-of-two scale; the maximum finite E2M1 value remains `6`. +This option applies UOS to Linear quantization. It does not implement the +paper's full MXAttention pipeline or Pre-Normalization Quantization (PNQ). + +Both use FP4 E2M1 with E8M0 scales, `axis=-1`, `block_size=32`, +`round_mode="rint"`. W4A8 always quantizes the activation to +`float8_e4m3fn` with `scale_alg=0` and the same grouping/rounding, without +UOS parameters. A4 and A8 share the prepared W4 and scale. UOS is currently +NPU-only; ROCm retains its existing AITER online algorithm. + +An explicit UOS online configuration is +`{"method": "mxfp4", "mxfp4_scale_alg": 2}`. For offline checkpoints the +setting affects **only A4 activations**. Existing weights are never requantized: +a UOS checkpoint can run UOS W4 + UOS A4 or UOS W4 + OCP A4, and both fall back to +UOS W4 + OCP A8. Record the weight-generation algorithm separately from the +runtime activation setting. Smooth remains active with either setting. The +UOS configuration explicitly selects `2`; attention computation is unchanged. + ### `mxfp4` — Single-Scale Online Mode -Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to -MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale -(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no -calibration `mul_scale` is available. Applies equally to single-transformer -and cascade (A14B) models — both transformers in a cascade receive the same -quantization config automatically. +Online mode quantizes BF16/FP16 weights once at loading and computes activation +scales during forward. Each block of 32 K elements shares one E8M0 scale; +no calibrated `mul_scale` is available. An explicit global quantization config +applies to both cascade experts. Per-component configs can select different +policies; checkpoint auto-detection resolves each expert independently. ```python from vllm_omni import Omni @@ -103,7 +289,7 @@ outputs = omni.generate( ```bash # Single-transformer or cascade model — same command -python text_to_video.py --model --quantization mxfp4 +python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4 # Online serving vllm serve --omni --quantization mxfp4 @@ -113,19 +299,19 @@ vllm serve --omni --quantization mxfp4 Online DualScale mode computes both fine and coarse scales on the fly from BF16 weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to -single-transformer and cascade (A14B) models. Compared to `mxfp4` online, -DualScale provides better quantization accuracy at higher compute cost. +single-transformer and cascade (A14B) models with compatible Linear layers. +Compare quality and performance using the intended model and workload. The default configuration keeps the leading 5 transformer blocks in BF16 -(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this -is sufficient to meet quality requirements and is the recommended setting. +(`num_bf16_fallback_layers=5`). This is an implementation default; select and +validate the BF16 policy against your own quality target. ```python omni = Omni(model="", quantization="mxfp4_dualscale") ``` ```bash -python text_to_video.py --model --quantization mxfp4_dualscale +python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4_dualscale ``` If accuracy debugging identifies additional precision-sensitive layers, they @@ -136,21 +322,22 @@ omni = Omni( model="", quantization_config={ "method": "mxfp4_dualscale", - "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override + "ignored_layers": ["blocks.10.attn1.to_qkv"], # explicit per-layer override }, ) ``` BF16 fallback routing in online mode applies two rules in priority order: -1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix - matches is kept in BF16 regardless of block index. +1. **`ignored_layers`** (explicit per-layer override): a listed runtime Linear + module is kept in BF16 regardless of block index. Use exact module paths, + not arbitrary block prefixes or names ending in `.weight`. 2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults - to `5` (recommended). Layers outside `blocks.N.*` - (e.g. `condition_embedder`) are always quantized. + to `5`. Layers outside `blocks.N.*` are quantized unless explicitly listed + in `ignored_layers`. -### `mxfp4_dualscale` — DualScale Offline Mode (Recommended) +### `mxfp4_dualscale` — DualScale Offline Mode Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A preprocessing step converts the raw quantized output to the diffusers format @@ -169,8 +356,8 @@ inside `config.json`. At runtime, each layer's prefix is matched against Each quantized linear layer stores four tensors: | Tensor | Shape | dtype | Description | -|--------|-------|-------|-------------| -| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) | +| -------- | ------- | ------- | ------------- | +| `weight` | `(N, K)` | float8_e4m3fn | Numeric FP4 values; packed into two values per byte after loading | | `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) | | `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale | | `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) | @@ -223,7 +410,7 @@ python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \ ``` | Argument | Description | -|----------|-------------| +| ---------- | ------------- | | `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` | | `--original-model` | Root directory of the original BF16 diffusers model | | `--quant-path` | Root directory of the msModelSlim quantized output | @@ -251,14 +438,14 @@ The `quantization_config` injected into `config.json` for each transformer: ``` `ignored_layers` lists every linear layer that retains its original BF16 weight, -using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out` +using vllm-omni runtime module paths (QKV-fused, FFN underscored, `to_out` unindexed). The exact entries are determined by the quantization tool (msModelSlim) and may differ between `transformer` and `transformer_2` in a cascade model. #### Step 3 — Serve ```bash -python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale +python examples/offline_inference/text_to_video/text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale # Online serving vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni @@ -275,25 +462,32 @@ omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale") ## Parameters -### `mxfp4` (single-scale, online only) +### `mxfp4` (single-scale, online + offline) | Parameter | Type | Default | Description | -|-----------|------|---------|-------------| +| ----------- | ------ | --------- | ------------- | | `method` | str | — | `"mxfp4"` | -| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 | +| `mxfp4_scale_alg` | int | `0` | MXFP4 block scaling: `0` OCP MX, `2` UOS (`Qmax=7.25`, NPU only). Online: W4 and A4; offline: A4 only. A8 remains OCP MXFP8 | +| `is_checkpoint_mxfp4_serialized` | bool | `false` | Load pre-quantized single-scale weights and scales; declared in checkpoint `config.json` | +| `require_smooth_scale` | bool | `false` | Require calibrated Smooth tensors for offline single-scale weights | +| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16 | +| `w4a8_fallback_layers` | list[str] | `[]` | Exact runtime Linear paths that always use A8; NPU only | +| `w4a8_fallback_steps` | list[int] | `[]` | Zero-based denoise steps that use A8; NPU only | ### `mxfp4_dualscale` (dual-scale, online + offline) | Parameter | Type | Default | Description | -|-----------|------|---------|-------------| +| ----------- | ------ | --------- | ------------- | | `method` | str | — | `"mxfp4_dualscale"` | | `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script | -| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override | -| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B | +| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override | +| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Validate the default of `5` against the intended model and quality target | +| `w4a8_fallback_layers` | list[str] | `[]` | Must remain empty; DualScale does not support W4A8 fallback | +| `w4a8_fallback_steps` | list[int] | `[]` | Must remain empty; DualScale does not support W4A8 fallback | #### BF16 fallback priority (online mode) -``` +```text for each linear layer: if prefix in ignored_layers → BF16 (explicit override, highest priority) elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule) @@ -307,8 +501,9 @@ unless they appear in `ignored_layers`. 1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using `npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is - available — all output partitions receive an identity pre-scale. No offline - checkpoint format exists for this method. + available — all output partitions receive an identity pre-scale. Offline + single-scale checkpoints use the canonical numeric FP4 layout + described above. 2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)** quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse @@ -316,7 +511,7 @@ unless they appear in `ignored_layers`. explicit `ignored_layers` stay in BF16 for accuracy. 3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** — - **recommended for production** — loads four tensors per quantized layer: FP4 + loads four tensors per quantized layer: FP4 weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale (`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback layers have no quantization tensors and are routed via `ignored_layers`. @@ -325,21 +520,21 @@ unless they appear in `ignored_layers`. layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy float32 round-trip. -5. **Cascade model config propagation**: in a cascade model (transformer + - transformer_2), vLLM-Omni reads each transformer's own `config.json` and - rebuilds the quant config locally when `ignored_layers` differs between - transformers, ensuring per-layer routing is accurate for each. The first - transformer's config is propagated to `od_config` so the second transformer - can reuse it as a starting point. +5. **Cascade model configuration**: checkpoint auto-detection resolves each + expert's own `config.json` independently; the first expert's detected config + is not reused as the second expert's policy. An explicit global runtime + config applies to both experts, while per-component configs select policies + by the exact `transformer` / `transformer_2` names. Serialized checkpoint + metadata still determines each expert's storage format and BF16 layers. 6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the fused name (`attn1.to_qkv`), written automatically by the merge script. -7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The - DualScale offline method mitigates this with calibrated `mul_scale` smooth - quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade - off compression vs. accuracy for precision-sensitive layers. +7. FP4 has fewer representable values than FP8, so evaluate quantization + error and generation quality for the chosen recipe. Offline Smooth scaling + and BF16 fallback can help with sensitive layers. + `num_bf16_fallback_layers` applies only to online DualScale. ## Adapting MXFP4 for a New Model @@ -350,27 +545,26 @@ a merge script for offline checkpoints. ### Step 1 — Discover runtime layer names -`ignored_layers` entries must match the **runtime parameter names** used inside -vllm-omni, which may differ from the names stored in the diffusers checkpoint. -The canonical source of truth is the model's own `named_parameters()`. +`ignored_layers` entries use exact **runtime Linear module paths**, without +`.weight`. Inspect the actual transformer instance in model initialization or +a developer harness; `Omni` does not expose an `omni.pipeline.transformer` API. +For a transformer instance named `transformer`: ```python -from vllm_omni import Omni +from vllm.model_executor.layers.linear import LinearBase -# Load the model without quantization to inspect parameter names. -omni = Omni(model="/path/to/your-model") # no --quantization flag -for name, _ in omni.pipeline.transformer.named_parameters(): - if "weight" in name and "scale" not in name: +for name, module in transformer.named_modules(): + if isinstance(module, LinearBase): print(name) ``` -Compare the printed names against the diffusers checkpoint keys +Compare these module paths against the corresponding diffusers checkpoint keys (`safetensors.safe_open` or `torch.load`) to identify any renames your model applies. Common patterns that differ in Wan2.2 (and may appear in other models): | Diffusers checkpoint name | vllm-omni runtime name | Reason | -|---------------------------|------------------------|--------| +| --------------------------- | ------------------------ | -------- | | `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` | | `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores | | `ffn.net.2` | `ffn.net_2` | Same underscore rule | @@ -382,8 +576,9 @@ fused runtime parameters. !!! warning "Partial QKV fallback is not allowed" If your model fuses Q, K, V into a single layer, `ignored_layers` must - include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but - `to_k`, `to_v` quantized) cannot be expressed at runtime because they share + name the fused runtime module (for Wan, `attn1.to_qkv`). The source + checkpoint must keep **all three or none** in BF16. A partial fallback + (e.g. `to_q` in BF16 but `to_k`, `to_v` quantized) cannot be expressed at runtime because they share one `QKVParallelLinear`. The merge script enforces this and raises an error if only some of the trio appear as non-quantized. @@ -392,12 +587,13 @@ fused runtime parameters. #### Online mode Pass `ignored_layers` directly in the quantization config using the **runtime -names** discovered in Step 1. No code changes to the model are required. +names** discovered in Step 1. This assumes the model already passes its +quantization config and module prefixes to compatible vLLM Linear layers. ```python omni = Omni( model="/path/to/your-model", - quantization={ + quantization_config={ "method": "mxfp4_dualscale", "ignored_layers": [ "blocks.0.attn1.to_qkv", # runtime name, not diffusers name @@ -409,9 +605,14 @@ omni = Omni( ``` ```bash -# CLI does not support list-typed ignored_layers directly. -# Use the Python API or set ignored_layers in config.json (offline). -python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale +# The text-to-video example accepts list-valued fields through JSON. +python examples/offline_inference/text_to_video/text_to_video.py \ + --model /path/to/your-model \ + --quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}' + +# Serving accepts the same configuration. +vllm serve /path/to/your-model --omni \ + --diffusion-quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}' ``` The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers @@ -436,9 +637,13 @@ the merge script is correct. The injected block: } ``` -To add a layer manually (e.g. to pin an additional layer to BF16 without -re-running the merge script), edit `config.json` inside the transformer -subfolder. Use runtime names, not diffusers checkpoint names. +Adding an offline BF16 fallback layer requires both the original BF16 weight +from the base model and a matching `ignored_layers` entry. Regenerate the +checkpoint with the revised storage policy, restoring the original weight +(and bias, if present) and removing that layer's quantization tensors. +Editing `config.json` alone does not reconstruct BF16 weights from stored FP4 +values. For fused QKV, restore all three projections together. Use runtime +module paths in `ignored_layers`. ### Step 3 — Write a merge script for offline mode @@ -474,6 +679,7 @@ renaming rules and collect the results. !!! tip "Validate before serving" After producing the offline checkpoint, load it without a `--quantization` flag and verify that vLLM-Omni auto-detects the correct method. Check that - the layer count reported in the startup log matches expectations: quantized - layer count + `ignored_layers` count should equal total linear layer count. - Any mismatch indicates a name-mapping bug in the merge script. + each intended runtime Linear module resolves to the expected quantized or + BF16 method and that every BF16 entry has original BF16 weights. Account + for fused modules and, in online DualScale, the leading-block fallback rule; + the raw length of `ignored_layers` alone is not a reliable layer count. diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 06179ef9e..7b7bcb890 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -384,13 +384,20 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) - parser.add_argument( + quantization_group = parser.add_mutually_exclusive_group() + quantization_group.add_argument( "--quantization", type=str, default=None, choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) + quantization_group.add_argument( + "--quantization-config", + type=json.loads, + default=None, + help='Quantization JSON, e.g. \'{"method":"mxfp4","w4a8_fallback_steps":[0,2]}\'.', + ) # Distributed and parallel execution parser.add_argument( @@ -572,6 +579,8 @@ def main(): omni_kwargs["flow_shift"] = args.flow_shift if args.quantization is not None: omni_kwargs["quantization"] = args.quantization + if args.quantization_config is not None: + omni_kwargs["quantization_config"] = args.quantization_config if args.cache_backend is not None: omni_kwargs["cache_backend"] = args.cache_backend omni_kwargs["cache_config"] = cache_config diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index e0a5ad9b0..1ed62f7f1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1480,7 +1480,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config - engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path)) od_config = OmniDiffusionConfig.from_kwargs(**engine_args) assert od_config.kv_transfer_config.engine_id == "dit-engine-1" @@ -1739,3 +1739,31 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder): builder(pipeline, deploy) else: builder(pipeline, user_deploy_config=deploy) + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit): + from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig + from vllm_omni.quantization import build_quant_config + + checkpoint = TransformerConfig.from_dict( + { + "quantization_config": { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "w4a8_fallback_steps": [37], + } + } + ) + requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + tf_model_config=checkpoint, + quantization_config=requested, + ) + assert cfg.quantization_config_is_auto_detected is not explicit + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint)) + cfg.enrich_config() + restored = ForkingPickler.loads(ForkingPickler.dumps(cfg)) + assert restored.quantization_config_is_auto_detected is not explicit + assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37]) diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py index bb95363fd..b01f275eb 100644 --- a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py +++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for transformer quant-config auto-detection. The loader path at pipeline_wan2_2.py carries two quantization contracts: @@ -27,6 +27,7 @@ from types import SimpleNamespace import pytest import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_ # od_config must remain unchanged — _create_transformer does not modify it. assert od_config.quantization_config is None + + +def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch): + import json + + import torch + + import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module + import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module + + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True) + for component in ("transformer", "transformer_2"): + directory = tmp_path / component + directory.mkdir() + config = dict(_MIN_CFG) + expert_quant_config = ( + disk_quant_config[0 if component == "transformer" else 1] + if isinstance(disk_quant_config, tuple) + else disk_quant_config + ) + if expert_quant_config is not None: + config["quantization_config"] = expert_quant_config + (directory / "config.json").write_text(json.dumps(config)) + (tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]})) + + class FakeComponent(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8) + + class FakeTransformer(FakeComponent): + def __init__(self, quant_config=None, **kwargs): + super().__init__() + self.received_quant_config = quant_config + + for module in (wan22_module, i2v_module): + monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent()) + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer) + pipeline_cls = { + "t2v": wan22_module.Wan22Pipeline, + "i2v": i2v_module.Wan22I2VPipeline, + "vace": vace_module.Wan22VACEPipeline, + }[pipeline_kind] + monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None) + od_config = OmniDiffusionConfig( + model=str(tmp_path), + model_class_name=pipeline_cls.__name__, + dtype=torch.bfloat16, + quantization_config=active, + boundary_ratio=0.875, + flow_shift=12.0, + enable_diffusion_pipeline_profiler=False, + ) + od_config.enrich_config() + pipeline = pipeline_cls(od_config=od_config) + return pipeline + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")]) +@pytest.mark.parametrize("offline", [False, True]) +def test_pipeline_constructor_routes_each_expert_quantization( + pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker +): + """Run the actual pipeline constructors, expert factories and quant-method dispatch.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.quantization import build_quant_config + + active = build_quant_config( + { + "transformer": { + "method": method, + "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"], + "w4a8_fallback_steps": [0, 1], + }, + "transformer_2": { + "method": method, + "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"], + "w4a8_fallback_steps": [], + }, + } + ) + disk_quant_config = None + if offline: + disk_quant_config = { + "quant_method": method, + flag: True, + "ignored_layers": ["blocks.12.ffn.net_2"], + "w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"], + "w4a8_fallback_steps": [9], + } + pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch) + for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)): + received = transformer.received_quant_config + assert received.get_name() == method + assert getattr(received, flag) == offline + assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else []) + selected = f"blocks.{10 + index}.attn1.to_qkv" + other = f"blocks.{11 - index}.attn1.to_qkv" + assert received.w4a8_fallback_layers == [selected] + assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer + assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer + if offline: + assert isinstance( + received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod + ) + assert pipeline.od_config.quantization_config is active + + +@pytest.mark.parametrize("explicit_none", [False, True]) +def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch): + from vllm_omni.quantization import build_quant_config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None} + active = build_quant_config(components) + pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active)) + pipeline._create_transformer(_MIN_CFG, component="transformer_2") + assert captured[-1].get("quant_config") is None + with pytest.raises(ValueError, match="disabled for component 'transformer_2'"): + pipeline._create_transformer( + {**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}}, + component="transformer_2", + ) + + +def test_pipeline_component_default_reaches_expert_factory(monkeypatch): + from vllm_omni.quantization import build_quant_config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}}) + pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active)) + pipeline._create_transformer(_MIN_CFG, component="transformer_2") + assert captured[-1]["quant_config"] is active.default_config + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize( + "configured_expert, with_default", + [ + ("transformer", False), + ("transformer", True), + ("transformer_2", False), + ("transformer_2", True), + ], +) +def test_pipeline_partial_expert_config_does_not_match_sibling( + pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch +): + from vllm_omni.quantization import build_quant_config + + component_spec = { + configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + } + if with_default: + component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]} + active = build_quant_config(component_spec) + pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch) + sibling = "transformer_2" if configured_expert == "transformer" else "transformer" + selected = getattr(pipeline, configured_expert).received_quant_config + other = getattr(pipeline, sibling).received_quant_config + assert selected is active.component_configs[configured_expert] + assert other is active.default_config + if with_default: + assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"] + else: + assert other is None + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize("different_ignored_layers", [False, True]) +def test_enriched_auto_quantization_preserves_each_expert_policy( + pipeline_kind, different_ignored_layers, tmp_path, monkeypatch +): + """Startup auto-detection must not turn the first expert into a user override.""" + high = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ["blocks.0.ffn.net_2"], + "w4a8_fallback_steps": [], + "w4a8_fallback_layers": [], + "mxfp4_scale_alg": 0, + "require_smooth_scale": False, + } + low = { + **high, + "ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"], + "w4a8_fallback_steps": [37], + "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"], + "mxfp4_scale_alg": 2, + "require_smooth_scale": True, + } + pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch) + assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert. + for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]: + config = transformer.received_quant_config + for policy in ( + "ignored_layers", + "w4a8_fallback_steps", + "w4a8_fallback_layers", + "mxfp4_scale_alg", + "require_smooth_scale", + ): + assert getattr(config, policy) == expected[policy], policy diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 526eeeb15..58a5e961b 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @pytest.fixture(autouse=True) -def _cpu_platform(monkeypatch): +def _cpu_pipeline_runtime(monkeypatch): + # These lightweight pipelines have no accelerator or distributed groups. module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2") monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False) + monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1)) class _StubTransformer(nn.Module): @@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: assert torch.equal(result, torch.full_like(latents, 10.0)) +@pytest.mark.parametrize("fail_first_request", [False, True]) +def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, + override_forward_context, + ) + from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step + + contexts: list[ForwardContext] = [] + events = [] + + class InjectedPredictionError(Exception): + pass + + class RecordingTransformer(_StubTransformer): + def __init__(self, expert): + super().__init__() + self.expert = expert + + def forward(self, hidden_states, encoder_hidden_states, **kwargs): + context = get_forward_context() + assert context is contexts[-1] + request_index = len(contexts) - 1 + step = context.denoise_step_idx + positive = bool(encoder_hidden_states[0, 0, 0] > 0) + events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2]))) + if fail_first_request and request_index == 0 and step == 1 and not positive: + raise InjectedPredictionError + return (torch.ones_like(hidden_states) * (1 if positive else -1),) + + def encode_prompt(**kwargs): + context = get_forward_context() + assert context.denoise_step_idx is None + contexts.append(context) + embeds = torch.ones(1, 8, 8) + return embeds, -embeds + + pipeline = _make_pipeline() + pipeline.transformer = RecordingTransformer("high") + pipeline.transformer_2 = RecordingTransformer("low") + pipeline.scheduler = _StubScheduler([900, 500, 100]) + monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt) + monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents) + + # Use the real runner context manager, Wan forward/diffuse, and CFG + # dispatch. Only model computation, scheduler math and unrelated KV I/O + # are replaced; no test code publishes a denoise step. + runner = object.__new__(DiffusionModelRunner) + runner.pipeline = pipeline + runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32) + runner.vllm_config = None + runner.cache_backend = None + monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None) + with override_forward_context(None): + for request_index in range(2): + request = OmniDiffusionRequest( + prompt="a fox walks", + request_id=f"request-{request_index}", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent" + ), + ) + + def execute(): + return runner._execute_request_list( + [request], + od_config=runner.od_config, + allow_single_output=True, + require_request_batch_support=False, + record_name="test_wan_step_context", + record_output_peak_memory=False, + ) + + if fail_first_request and request_index == 0: + with pytest.raises(InjectedPredictionError): + execute() + else: + execute() + assert not is_forward_context_available() + + assert len(contexts) == 2 and contexts[0] is not contexts[1] + expected = [ + (request, step, "high" if step == 0 else "low", positive, step in (0, 2)) + for request in range(2) + for step in range(2 if fail_first_request and request == 0 else 3) + for positive in (True, False) + ] + assert events == expected + + class _StubDMDScheduler: def __init__(self) -> None: self.predict_clean_calls: list[tuple[float, float, float]] = [] diff --git a/tests/diffusion/quantization/test_mxfp4_checkpoint.py b/tests/diffusion/quantization/test_mxfp4_checkpoint.py new file mode 100644 index 000000000..e183df078 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_checkpoint.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract.""" + +import sys +from types import ModuleType + +import pytest +import torch +from vllm.config.load import LoadConfig +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, +) +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleLinearMethod, + NPUMxfp4LinearMethod, + NPUMxfp4OnlineLinearMethod, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1) + + +class _CheckpointLayer(torch.nn.Module): + def __init__(self, method, dtype=torch.bfloat16): + super().__init__() + method.create_weights( + self, + input_size_per_partition=512, + output_partition_sizes=[2], + input_size=512, + output_size=2, + params_dtype=dtype, + weight_loader=default_weight_loader, + ) + + def load_weights(self, weights): + params = dict(self.named_parameters()) + loaded = set() + for name, value in weights: + if name in params: + params[name].weight_loader(params[name], value) + loaded.add(name) + return loaded + + +def _load_checkpoint(layer, config, weights, monkeypatch): + od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.counter_before_loading_weights = 0.0 + monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items())) + loader.load_weights(layer) + + +@pytest.mark.parametrize("serialized", [False, True]) +def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized) + # Exercise the shared implementation directly as well as the online MRO + # below: checkpoint requirements must follow the configured storage mode. + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)} + if serialized: + with pytest.raises(ValueError, match="Required weights.*weight_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch): + config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg) + method = NPUMxfp4OnlineLinearMethod(config) + layer = _CheckpointLayer(method) + assert set(dict(layer.named_parameters())) == {"weight"} + assert layer.weight.is_meta + + weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512) + packed_weight = torch.zeros((2, 256), dtype=torch.uint8) + generated_scale = torch.full((2, 16), 127, dtype=torch.uint8) + quantized_inputs = [] + + def quantize(value, **kwargs): + quantized_inputs.append(value.clone()) + expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg) + if scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert kwargs == expected + return packed_weight, generated_scale + + # Keep real creation, lazy loading, processing, and strict loader checks; + # substitute only the device quantization operator for this CPU test. + npu = ModuleType("torch_npu") + npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize) + monkeypatch.setitem(sys.modules, "torch_npu", npu) + _load_checkpoint(layer, config, {"weight": weight}, monkeypatch) + + assert len(quantized_inputs) == 1 + torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0) + torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0) + torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0) + assert not getattr(layer.weight_scale, "is_checkpoint_required", False) + method.process_weights_after_loading(layer) + assert len(quantized_inputs) == 1 + + +@pytest.mark.parametrize("require_smooth", [False, True]) +@pytest.mark.parametrize("has_smooth", [False, True]) +def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + scale = torch.linspace(0.5, 3.0, 512) + if has_smooth: + weights["mul_scale"] = scale + else: + weights.pop("mul_scale") + + if require_smooth and not has_smooth: + with pytest.raises(ValueError, match="mul_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + _load_checkpoint(layer, config, weights, monkeypatch) + torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0) + + +@pytest.mark.parametrize("bad_scale", ["shape", "name"]) +def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + if bad_scale == "shape": + weights["mul_scale"] = torch.ones(511) + with pytest.raises(AssertionError): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + weights["multiply_scale"] = weights.pop("mul_scale") + with pytest.raises(ValueError, match="mul_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + weights["mul_scale"][7] = bad_value + with pytest.raises(ValueError, match="finite, strictly positive"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)]) +def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + weights["mul_scale"] = bad_scale + with pytest.raises(ValueError, match="one-dimensional floating-point"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("bad_value", [1e10, 1e-10]) +def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + method = NPUMxfp4LinearMethod(config) + layer = _CheckpointLayer(method, torch.float16) + layer.mul_scale.data.fill_(bad_value) + # Validation must fail before any device packing or GEMM. + monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) + with pytest.raises(ValueError, match="finite, strictly positive"): + method.process_weights_after_loading(layer) + + +@pytest.mark.parametrize( + ("mode", "missing"), + [ + ("single", "weight"), + ("single", "weight_scale"), + ("dual", "weight"), + ("dual", "weight_scale"), + ("dual", "weight_dual_scale"), + ("dual", "mul_scale"), + ], +) +def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch): + if mode == "single": + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + method = NPUMxfp4LinearMethod(config) + else: + config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True) + method = NPUMxfp4DualScaleLinearMethod(config) + layer = _CheckpointLayer(method) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing} + with pytest.raises(ValueError, match=missing): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_single_scale_smooth_uses_real_row_parallel_loader(rank): + from vllm.model_executor.layers.linear import RowParallelLinear + + layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True))) + layer.tp_rank = rank + checkpoint_scale = torch.arange(1024, dtype=torch.float32) + RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale) + torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("fallback", [False, True]) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_single_scale_smooth_precedes_quantization_and_preserves_output( + dtype, fallback, with_bias, scale_alg, monkeypatch +): + config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg + ) + method = NPUMxfp4LinearMethod(config) + layer = _CheckpointLayer(method, dtype) + weights = { + "weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512), + "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8), + "mul_scale": torch.linspace(0.5, 3.0, 512), + } + _load_checkpoint(layer, config, weights, monkeypatch) + x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512) + bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None + quantized_inputs = [] + + # Only the two device operators are substituted. Exercise real apply(), + # step selection, Smooth multiplication, bias handling, and reshape logic. + npu = ModuleType("torch_npu") + + def quantize(value, **kwargs): + quantized_inputs.append(value.clone()) + expected = dict( + dst_type=torch.float8_e4m3fn if fallback else "fp4", + axis=-1, + block_size=32, + round_mode="rint", + scale_alg=0 if fallback else scale_alg, + ) + if not fallback and scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert kwargs == expected + return value, torch.ones((6, 16), dtype=torch.uint8) + + def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs): + if fallback: + assert output_dtype == torch.bfloat16 + if bias is not None: + assert bias.dtype == output_dtype + assert bias.shape == (1, 2) + result = x_q.float() @ weight.float() + if bias is not None: + result += bias.float() + return result.to(output_dtype) + + npu.__dict__.update( + float8_e4m3fn=torch.float8_e4m3fn, + float8_e8m0fnu=torch.uint8, + float4_e2m1fn_x2="fp4", + npu_dynamic_mx_quant=quantize, + npu_quant_matmul=matmul, + ) + monkeypatch.setitem(sys.modules, "torch_npu", npu) + with override_forward_context(ForwardContext()): + set_forward_context_denoise_step_idx(0 if fallback else 1) + result = method.apply(layer, x, bias) + + expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype) + assert len(quantized_inputs) == 1 + torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0) + expected = expected_input.float() @ weights["weight"].float().T + if bias is not None: + expected += bias.float() + if fallback: + expected = expected.bfloat16() + torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0) + assert result.dtype == dtype diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py index 240ee2509..c51b2013b 100644 --- a/tests/diffusion/quantization/test_mxfp4_config.py +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config.""" import pytest @@ -27,12 +27,42 @@ def test_mxfp4_config_get_name(): assert DiffusionMXFP4Config.get_name() == "mxfp4" +@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3]) +def test_mxfp4_scale_alg_rejects_invalid_values(value): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="mxfp4_scale_alg"): + build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value}) + + +@pytest.mark.parametrize("runtime_alg", [0, 2]) +@pytest.mark.parametrize("serialized", [False, True]) +def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg) + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "mxfp4_scale_alg": 2, + "ignored_layers": ["proj_out"], + }, + ) + assert resolved.is_checkpoint_mxfp4_serialized + assert resolved.mxfp4_scale_alg == runtime_alg + assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == [] + + def test_mxfp4_config_from_config_defaults(): from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config cfg = DiffusionMXFP4Config.from_config({}) assert cfg.is_checkpoint_mxfp4_serialized is False assert cfg.ignored_layers == [] + assert cfg.w4a8_fallback_layers == [] def test_mxfp4_config_from_config_serialized(): @@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized(): assert cfg.is_checkpoint_mxfp4_serialized is True +def test_mxfp4_config_requires_offline_smooth(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True} + assert DiffusionMXFP4Config.from_config(values).require_smooth_scale + assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale + with pytest.raises(ValueError, match="offline"): + build_quant_config({"method": "mxfp4", "require_smooth_scale": True}) + with pytest.raises(ValueError, match="boolean"): + DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"}) + + +@pytest.mark.parametrize("serialized", [False, True]) +def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1]) + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "require_smooth_scale": True, + }, + ) + assert resolved.require_smooth_scale + assert resolved.is_checkpoint_mxfp4_serialized + assert resolved.w4a8_fallback_steps == [1] + + +@pytest.mark.parametrize("disk_requirement", [None, False, True]) +def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + disk = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ["blocks.0.attn2.to_q"], + } + if disk_requirement is not None: + disk["require_smooth_scale"] = disk_requirement + resolved = resolve_quant_config_from_disk(active, disk) + assert resolved.require_smooth_scale + assert resolved.ignored_layers == disk["ignored_layers"] + + +@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}]) +def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config( + ignored_layers=["blocks.0.attn2.to_q"], + w4a8_fallback_steps=[1], + w4a8_fallback_layers=["blocks.1.attn1.to_qkv"], + ) + + resolved = resolve_quant_config_from_disk(active, disk_qc) + + assert resolved is active + assert resolved.ignored_layers == ["blocks.0.attn2.to_q"] + assert resolved.w4a8_fallback_steps == [1] + assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"] + + +def test_quantization_disk_string_rejects_mismatched_active_method(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = build_quant_config("mxfp4") + + with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"): + resolve_quant_config_from_disk(active, "mxfp4_dualscale") + + +def test_quantization_disk_string_accepts_equivalent_method_alias(): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + class ActiveAliasConfig: + @staticmethod + def get_name() -> str: + return "inc" + + active = ActiveAliasConfig() + + assert resolve_quant_config_from_disk(active, "auto-round") is active + + def test_mxfp4_config_from_config_ignored_layers(): from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config @@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback(): assert cfg.ignored_layers == ["proj_out"] +@pytest.mark.parametrize("serialized", [False, True]) +@pytest.mark.parametrize( + "policy", + [ + {"w4a8_fallback_steps": [0]}, + {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + {"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + ], +) +def test_dualscale_rejects_w4a8_configuration(serialized, policy): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy} + with pytest.raises(ValueError, match="does not support W4A8"): + DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy) + with pytest.raises(ValueError, match="does not support W4A8"): + DiffusionMXFP4DualScaleMixedConfig.from_config(config) + with pytest.raises(ValueError, match="does not support W4A8"): + build_quant_config({"transformer": config}) + with pytest.raises(ValueError, match="does not support W4A8"): + resolve_quant_config_from_disk(None, config) + + +@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"]) +@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)]) +def test_mxfp4_rejects_invalid_fallback_steps(method, steps): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="non-negative integer"): + build_quant_config({"method": method, "w4a8_fallback_steps": steps}) + + +def test_mxfp4_fallback_steps_follow_request_context(): + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step + + steps = [2, 0, 2] + cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps}) + steps.append(1) + assert cfg.w4a8_fallback_steps == [0, 2] + with override_forward_context(None): + assert not _is_w4a8_fallback_step([]) + with pytest.raises(RuntimeError, match="publishes its denoise step"): + _is_w4a8_fallback_step(cfg.w4a8_fallback_steps) + for _request in range(2): + with override_forward_context(ForwardContext()): + for step, expected in [(0, True), (1, False), (2, True), (3, False)]: + set_forward_context_denoise_step_idx(step) + # Both CFG passes and both experts read the same index. + assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [ + expected, + expected, + ] + with pytest.raises(RuntimeError, match="publishes its denoise step"): + _is_w4a8_fallback_step(cfg.w4a8_fallback_steps) + + +@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")]) +@pytest.mark.parametrize("steps", [[], [0, 2]]) +@pytest.mark.parametrize("already_serialized", [False, True]) +def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + layers = ["blocks.10.attn1.to_qkv"] if steps else [] + active = build_quant_config( + { + "method": method, + flag: already_serialized, + "w4a8_fallback_steps": steps, + "w4a8_fallback_layers": layers, + } + ) + # Separate experts may have distinct high precision layers. Both must + # retain the requested step policy during online/offline reconciliation. + for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]: + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": method, + flag: True, + "ignored_layers": ignored, + "w4a8_fallback_steps": [9], + "w4a8_fallback_layers": ["blocks.9.ffn.net_2"], + }, + ) + assert getattr(resolved, flag) + assert resolved.ignored_layers == ignored + assert resolved.w4a8_fallback_steps == steps + assert resolved.w4a8_fallback_layers == layers + assert active.ignored_layers == [] + + +@pytest.mark.parametrize( + "runtime_policy", + [{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}], +) +def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = build_quant_config({"method": "mxfp4", **runtime_policy}) + resolved_experts = [] + for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]): + disk = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ignored, + "w4a8_fallback_steps": [9], + "w4a8_fallback_layers": ["blocks.9.ffn.net_2"], + } + resolved = resolve_quant_config_from_disk(active, disk) + assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", []) + assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", []) + assert resolved.ignored_layers == ignored + detected = resolve_quant_config_from_disk(None, disk) + assert detected.w4a8_fallback_steps == [9] + assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"] + resolved_experts.append(resolved) + assert resolved_experts[0] is not resolved_experts[1] + assert active.is_checkpoint_mxfp4_serialized is False + assert active.ignored_layers == [] + + +@pytest.mark.parametrize( + "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")] +) +def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active_spec = { + "method": method, + flag: True, + "ignored_layers": ["blocks.0.attn2.to_q"], + "w4a8_fallback_steps": [1] if method == "mxfp4" else [], + "w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [], + } + if method == "mxfp4": + active_spec["require_smooth_scale"] = True + active = build_quant_config(active_spec) + + resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True}) + + assert resolved is not active + assert resolved.ignored_layers == [] + assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps + assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers + if method == "mxfp4": + assert resolved.require_smooth_scale + + +@pytest.mark.parametrize( + "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")] +) +@pytest.mark.parametrize("with_active", [False, True]) +def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = None + disk_steps = [9] if method == "mxfp4" else [] + active_steps = [1] if method == "mxfp4" else [] + expected_steps = disk_steps + if with_active: + active = build_quant_config( + { + "method": method, + flag: True, + "ignored_layers": ["blocks.0.attn2.to_q"], + "w4a8_fallback_steps": active_steps, + } + ) + expected_steps = active_steps + + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": method, + flag: True, + "modules_to_not_convert": ["blocks.2.ffn.net_2"], + "w4a8_fallback_steps": disk_steps, + }, + ) + + assert resolved.ignored_layers == ["blocks.2.ffn.net_2"] + assert resolved.w4a8_fallback_steps == expected_steps + + # --------------------------------------------------------------------------- # build_quant_config integration # --------------------------------------------------------------------------- @@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch): assert isinstance(method, ROCmMxfp4OnlineLinearMethod) -def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch): +@pytest.mark.parametrize("steps", [[], [0, 2]]) +def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps): """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe.""" from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") - cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"]) + cfg = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps + ) assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod) +def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2]) + + with pytest.raises(NotImplementedError, match="only supported on NPU"): + cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + + def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch): """MXFP4 on ROCm requires gfx950; any other arch must raise.""" from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config @@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform): assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) assert layer.input_size_per_partition == _TP2_K // _TP2 assert layer.output_size_per_partition == _TP2_N + + +@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"]) +@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]]) +def test_mxfp4_rejects_invalid_fallback_layers(method, layers): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="exact runtime Linear paths"): + build_quant_config({"method": method, "w4a8_fallback_layers": layers}) + + +@pytest.mark.parametrize("method", ["mxfp4"]) +def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.component_config import resolve_component_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"] + components = build_quant_config( + {"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}} + ) + names.append("blocks.11.attn1.to_qkv") + assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]] + assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == [] + detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]}) + assert detected.w4a8_fallback_layers == [names[0]] + assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]] + + +@pytest.mark.parametrize( + "kind, serialized_flag", + [("mxfp4", "is_checkpoint_mxfp4_serialized")], +) +@pytest.mark.parametrize("offline", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_mxfp4_layer_and_step_union_routes_actual_methods( + kind, serialized_flag, offline, scale_alg, monkeypatch, mocker +): + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.quantization import build_quant_config, mxfp4_config + + monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True) + selected = "blocks.10.attn1.to_qkv" + other = "blocks.11.attn2.to_q" + preserved = "blocks.12.ffn.net_2" + config = build_quant_config( + { + "method": kind, + "mxfp4_scale_alg": scale_alg, + serialized_flag: offline, + "w4a8_fallback_layers": [selected, preserved], + "ignored_layers": [preserved], + "w4a8_fallback_steps": [1], + } + ) + # Real config dispatch creates separate methods for selected and other layers. + linear = mocker.Mock(spec=LinearBase) + assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod) + x = torch.ones(2, 64, dtype=torch.bfloat16) + layer = torch.nn.Module() + for name in ("weight", "weight_scale", "mul_scale"): + layer.register_buffer(name, torch.ones(64, dtype=x.dtype)) + calls = [] + + def record_a8(*args): + calls.append("A8") + return x + + def record_a4(*args): + calls.append("A4") + return x + + monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8) + for name in (selected, other): + method = config.get_quant_method(linear, name) + assert method.is_w4a8_fallback_layer == (name == selected) + scales = (x, x) + monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales) + monkeypatch.setattr(method, "_quant_matmul", record_a4) + for step in (0, 1, 2): + with override_forward_context(ForwardContext(denoise_step_idx=step)): + method.apply(layer, x) + assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"] + # A permanently selected layer works without a request context, even if steps are also configured. + with override_forward_context(None): + config.get_quant_method(linear, selected).apply(layer, x) + assert calls[-1] == "A8" + + +@pytest.mark.parametrize("online", [False, True]) +@pytest.mark.parametrize("selected", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch): + import sys + from types import SimpleNamespace + + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False) + prepared = [] + packed = [] + + def quantize(x, **kwargs): + prepared.append(kwargs) + return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8) + + def pack(x, dtype): + packed.append(x) + return x.to(dtype) + + fake_npu = SimpleNamespace( + float4_e2m1fn_x2=torch.int8, + float8_e8m0fnu=torch.uint8, + float8_e4m3fn=torch.float8_e4m3fn, + npu_dtype_cast=pack, + npu_dynamic_mx_quant=quantize, + ) + monkeypatch.setitem(sys.modules, "torch_npu", fake_npu) + config = mxfp4_config.DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=not online, + w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [], + w4a8_fallback_steps=[1], + mxfp4_scale_alg=scale_alg, + ) + cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod + method = cls(config, prefix="blocks.0.attn1.to_qkv") + layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16) + layer.orig_dtype = torch.bfloat16 + if not online: + layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False)) + layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False)) + method.process_weights_after_loading(layer) + pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + method.process_weights_after_loading(layer) + assert len(prepared) == int(online) + expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg) + if scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert prepared == ([expected] if online else []) + assert len(packed) == int(not online) + branches = [] + + def a4(x_q, x_scale, target, bias, dtype): + assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers + branches.append("A4") + return torch.zeros(x_q.shape[0], 2, dtype=dtype) + + def a8(x, weight, scale, bias, dtype): + assert (weight.data_ptr(), scale.data_ptr()) == pointers + branches.append("A8") + return torch.zeros(x.shape[0], 2, dtype=dtype) + + monkeypatch.setattr(method, "_quant_matmul", a4) + monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8) + monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x)) + monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization")) + monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking")) + for step in (0, 1, 2): + with override_forward_context(ForwardContext(denoise_step_idx=step)): + method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16)) + assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"]) + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers + assert set(layer.state_dict()) == ( + {"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"} + ) + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + + +@pytest.mark.parametrize("online", [False, True]) +def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch): + import sys + from types import SimpleNamespace + + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False) + original_ones = torch.ones + monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"})) + calls = [] + + def quantize_dual(x, **kw): + calls.append(kw.get("smooth_scale")) + return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8) + + fake_npu = SimpleNamespace( + float4_e2m1fn_x2=torch.int8, + npu_dtype_cast=lambda x, dtype: x.to(dtype), + npu_format_cast=lambda x, *args, **kw: x, + npu_dynamic_dual_level_mx_quant=quantize_dual, + ) + monkeypatch.setitem(sys.modules, "torch_npu", fake_npu) + config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[] + ) + cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod + method = cls(config) + layer = torch.nn.Module() + layer.input_size_per_partition = 512 + for name, value in { + "weight": torch.ones(2, 512, dtype=torch.bfloat16), + "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8), + "weight_dual_scale": torch.ones(2, 1, 1), + "mul_scale": torch.ones(512), + }.items(): + layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False)) + method.process_weights_after_loading(layer) + method.process_weights_after_loading(layer) + assert len(calls) == int(online) + assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"} + assert not hasattr(layer, "w4a8_weight") + monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16)) + method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16)) + assert len(calls) == 1 + int(online) + torch.testing.assert_close(calls[-1], layer.mul_scale) diff --git a/tests/diffusion/quantization/test_wan_mxfp4_loading.py b/tests/diffusion/quantization/test_wan_mxfp4_loading.py new file mode 100644 index 000000000..fc7a43d43 --- /dev/null +++ b/tests/diffusion/quantization/test_wan_mxfp4_loading.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU.""" + +import pytest +import torch +from vllm.config.load import LoadConfig +from vllm.model_executor.layers.linear import QKVParallelLinear + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True) + prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer" + monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1) + + +class _WanCheckpoint(torch.nn.Module): + load_weights = WanTransformer3DModel.load_weights + _log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary + + def __init__(self, config, *, num_blocks=1): + super().__init__() + blocks = [] + for block_idx in range(num_blocks): + block = torch.nn.Module() + block.attn1 = torch.nn.Module() + block.attn1.to_qkv = QKVParallelLinear( + hidden_size=512, + head_size=32, + total_num_heads=4, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix=f"blocks.{block_idx}.attn1.to_qkv", + disable_tp=True, + ) + blocks.append(block) + self.blocks = torch.nn.ModuleList(blocks) + + +def _checkpoint(mode, fused): + if mode == "single": + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + elif mode == "dual": + config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True) + else: + config = None + model = _WanCheckpoint(config) + weights = {} + expected = {} + for name, param in model.named_parameters(): + suffix = name.rsplit(".", 1)[1] + output_dim = getattr(param, "output_dim", None) + if output_dim is None: + value = torch.full_like(param, 2.5) + shards = [value] * 3 + else: + parts = param.chunk(3, dim=output_dim) + shards = [ + torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts) + ] + value = torch.cat(shards, dim=output_dim) + expected[name] = value + if fused: + weights[name] = value + else: + for shard_id, shard in zip(("q", "k", "v"), shards): + weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard + return model, config, weights, expected + + +def _load(model, config, weights, monkeypatch): + loader = DiffusersPipelineLoader( + LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config) + ) + loader.counter_before_loading_weights = 0.0 + monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items())) + loader.load_weights(model) + + +@pytest.mark.parametrize("mode", ["single", "dual", "bf16"]) +@pytest.mark.parametrize("fused", [False, True]) +def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch): + model, config, weights, expected = _checkpoint(mode, fused) + _load(model, config, weights, monkeypatch) + for name, param in model.named_parameters(): + torch.testing.assert_close(param, expected[name], rtol=0, atol=0) + + +@pytest.mark.parametrize("mode", ["single", "dual", "bf16"]) +def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights.pop("blocks.0.attn1.to_k.weight") + weights.pop("blocks.0.attn1.to_v.weight") + with pytest.raises(ValueError, match="to_qkv.weight"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("mode", ["single", "dual"]) +def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights.pop("blocks.0.attn1.to_k.weight_scale") + weights.pop("blocks.0.attn1.to_v.weight_scale") + with pytest.raises(ValueError, match="to_qkv.weight_scale"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("mode", ["single", "dual"]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2 + if reverse: + weights = dict(reversed(weights.items())) + with pytest.raises(ValueError, match="must share the same Smooth tensor"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"]) +def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name): + model, _, _, _ = _checkpoint("single", fused=True) + loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))]) + assert loaded == set() + + +@pytest.mark.parametrize("fused", [False, True]) +def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch): + model, config, weights, _ = _checkpoint("single", fused=fused) + projection = "to_qkv" if fused else "to_q" + weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1) + with pytest.raises(ValueError, match="cannot load DualScale tensor"): + _load(model, config, weights, monkeypatch) + + +def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override(): + config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"]) + model = _WanCheckpoint(config) + # Uses the same validation called at the end of WanTransformer3DModel construction. + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"): + config.w4a8_fallback_layers = [bad] + with pytest.raises(ValueError, match="exact runtime Linear paths"): + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + + +def test_wan_layer_policy_defers_only_other_pp_stage_paths(): + from vllm.model_executor.models.utils import PPMissingLayer + + config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"]) + model = _WanCheckpoint(config) + model.blocks.append(PPMissingLayer()) + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + config.w4a8_fallback_layers = ["blocks.1"] + with pytest.raises(ValueError, match="exact runtime Linear paths"): + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + + +@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"]) +def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker): + selected = "blocks.0.attn1.to_qkv" + bf16 = "blocks.1.attn1.to_qkv" + transformer_config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=is_checkpoint_serialized, + w4a8_fallback_layers=[selected], + ) + transformer_2_config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=is_checkpoint_serialized, + ignored_layers=[bf16], + w4a8_fallback_layers=[selected, bf16], + ) + pipeline = torch.nn.Module() + pipeline.transformer = _WanCheckpoint(transformer_config) + pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2) + WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config) + WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config) + + loader = DiffusersPipelineLoader( + LoadConfig(), + OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config), + ) + events = [] + summaries = [] + + def process_weights(model, _target_device): + events.append("process") + for module in model.modules(): + method = getattr(module, "quant_method", None) + if getattr(method, "is_w4a8_fallback_layer", False): + module._already_called_process_weights_after_loading = True + + def capture_summary(message, *args): + rendered = message % args + summaries.append(rendered) + events.append(f"summary:{args[0]}") + + loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign] + loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = process_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + mocker.patch( + "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info", + side_effect=capture_summary, + ) + + assert loader.load_model(load_device="cpu") is pipeline + assert events == ["load", "process", "summary:transformer", "summary:transformer_2"] + assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0] + assert "BF16 overrides=0 []" in summaries[0] + assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0] + assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1] + assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1] + assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1] + + +def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker): + selected = "blocks.0.attn1.to_qkv" + config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected]) + model = _WanCheckpoint(config) + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info") + + WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer") + + message, *args = logger_info.call_args.args + rendered = message % tuple(args) + assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 23ff796e2..6f37f823d 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.fixture(autouse=True) +def _local_model_paths(monkeypatch): + # These tests exercise config transport, not model repository resolution. + monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model) + + def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: """Simulate the real path: create_default_diffusion → OmniDiffusionConfig. diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py index c588b7a38..e0c381bf3 100644 --- a/tests/examples/offline_inference/test_text_to_video_prompts.py +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import functools import importlib.util +import sys +from types import SimpleNamespace import pytest import torch @@ -24,6 +26,51 @@ def _load_text_to_video(): return module +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + (["--quantization", "mxfp4"], {"quantization": "mxfp4"}), + ( + ["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'], + {"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}}, + ), + ], +) +def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch): + mod = _load_text_to_video() + captured = {} + + class StopBeforeModelLoadError(Exception): + pass + + def capture_omni(**kwargs): + captured.update(kwargs) + raise StopBeforeModelLoadError + + monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments]) + monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu")) + monkeypatch.setattr(mod, "Omni", capture_omni) + + with pytest.raises(StopBeforeModelLoadError): + mod.main() + + assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected + + +def test_quantization_cli_options_are_mutually_exclusive(monkeypatch): + mod = _load_text_to_video() + monkeypatch.setattr( + sys, + "argv", + ["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'], + ) + + with pytest.raises(SystemExit) as exc: + mod.parse_args() + + assert exc.value.code == 2 + + @pytest.mark.parametrize( ("negative_prompt", "expected_negative_prompt"), [(None, None), ("", ""), ("blurry", "blurry")], diff --git a/tests/platforms/npu/quant/test_mxfp4_step_fallback.py b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py new file mode 100644 index 000000000..8431ebf17 --- /dev/null +++ b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights.""" + +import pytest +import torch + +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu] + +torch_npu = pytest.importorskip("torch_npu") + + +@torch.inference_mode() +def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2): + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + NPUMxfp4LinearMethod, + NPUMxfp4OnlineLinearMethod, + ) + + prefix = "blocks.10.attn1.to_qkv" + steps = [1] if policy in ("step", "layer_and_step") else [] + layers = [prefix] if policy in ("layer", "layer_and_step") else [] + fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg} + torch.manual_seed(42) + layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype) + if mode == "single": + method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix) + elif mode == "single_offline": + method = NPUMxfp4LinearMethod( + DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix + ) + layer.orig_dtype = dtype + values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0]) + layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight)) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False) + ) + layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False)) + else: + raise ValueError(f"Unsupported single-scale test mode: {mode}") + method.process_weights_after_loading(layer) + x = torch.randn(input_shape, device="npu", dtype=dtype) + pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + method.process_weights_after_loading(layer) + quantize = torch_npu.npu_dynamic_mx_quant + activation_types = [] + + def record_quantize(x, **kwargs): + dst = kwargs["dst_type"] + expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg + assert kwargs["scale_alg"] == expected_alg + assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint" + if expected_alg == 2: + assert kwargs["dst_type_max"] == 7.25 + else: + assert "dst_type_max" not in kwargs + activation_types.append(dst) + return quantize(x, **kwargs) + + monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize) + results = [] + for _request in range(2): + with override_forward_context(ForwardContext()): + for step in range(3): + set_forward_context_denoise_step_idx(step) + output = method.apply(layer, x, layer.bias) + torch.npu.synchronize() + assert output.shape == (*input_shape[:-1], 128) + assert output.dtype == dtype + assert torch.isfinite(output).all() + results.append(output.cpu()) + expected_types = [ + torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3) + ] + assert activation_types == expected_types * 2 + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + for first, second in zip(results[:3], results[3:]): + torch.testing.assert_close(first, second, rtol=0, atol=0) + + +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)]) +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch): + _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step") + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_single_scale_step_fallback_matches_independent_dense_reference(dtype): + if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name(): + pytest.skip("MXFP4 kernels require an Ascend950 NPU") + + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + NPUMxfp4LinearMethod, + ) + + # Synthetic serialized single-scale values and E8M0 exponent bytes. + # Powers of two make both activation precisions exact: this independent + # CPU dense oracle needs no production packing, dequantization or apply. + # Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate. + values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512) + weight = torch.cat((values, values), dim=1) + fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8) + smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0))) + bias = torch.tensor([0.25, -0.25]).repeat(64) + x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64 + expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias) + + layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype) + layer.weight.copy_(weight) + layer.bias.copy_(bias) + layer.orig_dtype = dtype + for name, value in {"weight_scale": fine, "mul_scale": smooth}.items(): + layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False)) + method = NPUMxfp4LinearMethod( + DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2 + ) + ) + method.process_weights_after_loading(layer) + prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"} + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + x = x.to(device="npu", dtype=dtype) + + with override_forward_context(ForwardContext()): + for step in range(3): + set_forward_context_denoise_step_idx(step) + actual = method.apply(layer, x, layer.bias).cpu() + # A8's documented GEMM output is BF16 before restoring the caller + # dtype; the reference does not reuse quantization or unpacking. + reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + assert actual.dtype == dtype + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +@pytest.mark.parametrize("policy", ["layer", "layer_and_step"]) +def test_mxfp4_layer_fallback(mode, policy, monkeypatch): + """Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix.""" + _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy) + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +def test_mxfp4_ocp_compatibility(mode, monkeypatch): + """OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests.""" + _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 56b40b836..c13935489 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -817,6 +817,8 @@ class _DiffusionConfigProjection: enable_stage_verification: bool = True prompt_file_path: str | None = None quantization_config: _QuantizationConfigType = None + # Internal provenance, retained across config projection and worker transport. + quantization_config_is_auto_detected: bool = False extras: dict[str, Any] = field(default_factory=dict) @field_validator("kv_transfer_config", mode="before") @@ -964,6 +966,8 @@ class _DiffusionConfigProjection: or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) ) if should_use_checkpoint_config: + if self.quantization_config is None: + self.quantization_config_is_auto_detected = True self.quantization_config = quant_config if quant_method is not None: self.additional_config.setdefault("auto_detected_quant_method", quant_method) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e1ce550d4..0ffd0c54e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1021,6 +1021,8 @@ class OmniDiffusionConfig: # str is resolved to {"method": } internally. # Per-component: {"transformer": {"method": "fp8"}, "vae": None} quantization_config: str | QuantizationConfig | dict[str, Any] | None = None + # Internal provenance, retained across config projection and worker transport. + quantization_config_is_auto_detected: bool = False # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint # has already resolved to vLLM's ModelOpt FP8 linear method. @@ -1358,6 +1360,8 @@ class OmniDiffusionConfig: or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) ) if should_use_checkpoint_config: + if self.quantization_config is None: + self.quantization_config_is_auto_detected = True self.quantization_config = tf_config.quant_config logger.info( "Auto-detected quantization '%s' from model config", diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 78247e7b9..834887f41 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -7,7 +7,7 @@ import json import os import re import time -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Sequence from pathlib import Path from typing import cast @@ -481,8 +481,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin): ) def _get_expected_parameter_names(self, model: nn.Module) -> set[str]: - """Return parameter names that should be covered by strict load checks.""" - all_parameter_names = {name for name, _ in model.named_parameters()} + """Return parameter names that should be covered by strict load checks. + + A parameter with an initialized checkpoint default can explicitly set + ``is_checkpoint_optional``. It is still loaded when present on disk. + """ + all_parameter_names = { + name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False) + } sources = self._get_weight_sources(model) # Keep strict behavior if no source metadata exists. @@ -853,9 +859,24 @@ class DiffusersPipelineLoader(HWRLoaderMixin): del model return self.load_fresh_canonical_model() raise + self._log_w4a8_fallback_load_summaries(model) self._attach_offload_startup_state(model) return model + @staticmethod + def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None: + """Ask discovered DiTs to report W4A8 state at the common load exit. + + Each DiT derives readiness from its processed layers, so a deferred + weight plan cannot be mistaken for the ordinary or HSDP load path. + """ + components = ModuleDiscovery.discover(model) + for component_name, dit in zip(components.dit_names, components.dits): + candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None) + if callable(candidate): + reporter = cast(Callable[[str], None], candidate) + reporter(component_name) + @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: """Ask online-quant layers to return to host memory once quantized. @@ -1065,6 +1086,16 @@ class DiffusersPipelineLoader(HWRLoaderMixin): # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights + # Offline formats can require scales that older online loaders + # were allowed to synthesize. Do not apply that legacy tolerance + # to an explicitly required checkpoint tensor. + required_missing = { + name + for name, param in model.named_parameters() + if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False) + } + if required_missing: + raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}") # NOTE: if the model is quantized, ignore not_loaded check for scale # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a # static activation `input_scale`, which the quant method may diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 15b036864..915fb45f7 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc return {} +def resolve_wan_transformer_quant_config( + config: dict, quant_config: QuantizationConfig | None, component: str +) -> QuantizationConfig | None: + """Resolve the expert before applying its checkpoint's storage contract.""" + from vllm_omni.quantization.component_config import ComponentQuantizationConfig + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + # Wan experts are siblings: "transformer_2" must not inherit "transformer" + # through the generic layer-prefix resolver. Select the whole expert name. + component_quant_config = ( + quant_config.component_configs.get(component, quant_config.default_config) + if isinstance(quant_config, ComponentQuantizationConfig) + else quant_config + ) + quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None + resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config")) + if quantization_disabled and resolved_quant_config is not None: + raise ValueError( + f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. " + "Use a BF16 checkpoint for this component or enable its matching quantization method." + ) + return resolved_quant_config + + def create_transformer_from_config( - config: dict, quant_config: QuantizationConfig | None = None, prefix: str = "" + config: dict, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + *, + component: str = "transformer", ) -> WanTransformer3DModel: """Create WanTransformer3DModel from config dict.""" kwargs: dict = {} @@ -216,11 +244,7 @@ def create_transformer_from_config( if "pos_embed_seq_len" in config: kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"] - if "quantization_config" in config: - from vllm_omni.quantization.factory import resolve_quant_config_from_disk - - quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"]) - + quant_config = resolve_wan_transformer_quant_config(config, quant_config, component) if quant_config is not None: kwargs["quant_config"] = quant_config if prefix: @@ -460,13 +484,13 @@ class Wan22Pipeline( # Initialize transformers with correct config (weights loaded via load_weights) if load_transformer: transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = self._create_transformer(transformer_config) + self.transformer = self._create_transformer(transformer_config, component="transformer") else: self.transformer = None if load_transformer_2: transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only) - self.transformer_2 = self._create_transformer(transformer_2_config) + self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2") else: self.transformer_2 = None @@ -504,10 +528,17 @@ class Wan22Pipeline( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) - def _create_transformer(self, config: dict) -> WanTransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel: """Create a transformer from a config dict. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_transformer_from_config( + config, + quant_config=quant_config, + component=component, + ) @property def guidance_scale(self): diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 01bd53bb8..1a75c40da 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -304,25 +304,10 @@ class Wan22I2VPipeline( # Transformers (weights loaded via load_weights) # Load config from model directory or HF Hub to get correct in_channels for I2V models transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = create_transformer_from_config( - transformer_config, - quant_config=od_config.quantization_config, - ) + self.transformer = self._create_transformer(transformer_config, component="transformer") if self.has_transformer_2: transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only) - t2_quant = transformer_2_config.get("quantization_config") - if isinstance(t2_quant, dict) and "quant_method" in t2_quant: - from vllm_omni.quantization.factory import build_quant_config - - method = t2_quant["quant_method"] - kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"} - t2_quant = build_quant_config(method, **kwargs) - else: - t2_quant = None - self.transformer_2 = create_transformer_from_config( - transformer_2_config, - quant_config=t2_quant, - ) + self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2") else: self.transformer_2 = None @@ -462,10 +447,13 @@ class Wan22I2VPipeline( image_embeds = self.image_encoder(pixel_values, output_hidden_states=True) return image_embeds.hidden_states[-2] - def _create_transformer(self, config: dict) -> WanTransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel: """Create a transformer from a config dict. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_transformer_from_config(config, quant_config=quant_config, component=component) def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: sampling_params_list = req.sampling_params_list diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index aeecff608..0d5f3cd67 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( Wan22Pipeline, load_wan_weights_with_optional_gate, + resolve_wan_transformer_quant_config, retrieve_latents, ) from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( @@ -60,6 +61,8 @@ def create_vace_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + component: str = "transformer", ) -> WanVACETransformer3DModel: """Create WanVACETransformer3DModel from config dict.""" kwargs = {} @@ -97,6 +100,7 @@ def create_vace_transformer_from_config( kwargs["vace_layers"] = config["vace_layers"] if "vace_in_channels" in config: kwargs["vace_in_channels"] = config["vace_in_channels"] + quant_config = resolve_wan_transformer_quant_config(config, quant_config, component) if quant_config is not None: kwargs["quant_config"] = quant_config if prefix: @@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): super().__init__(od_config=od_config, prefix=prefix) - def _create_transformer(self, config: dict) -> WanVACETransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel: """Build VACE transformer. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_vace_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_vace_transformer_from_config(config, quant_config=quant_config, component=component) def diffuse( self, diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index dbc555604..2808dcee7 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -18,7 +18,7 @@ from vllm.distributed import ( ) from vllm.logger import init_logger from vllm.model_executor.layers.conv import Conv3dLayer -from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import ( @@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod logger = init_logger(__name__) @@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module): # ROPE helper self._cached_rope_emb = None self._cached_rope_resolution = None + self._validate_w4a8_fallback_layers(quant_config) + + def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None: + """Reject checkpoint aliases/typos before loading weights or running inference.""" + requested = tuple(getattr(quant_config, "w4a8_fallback_layers", [])) + self._w4a8_fallback_layers = requested + self._local_w4a8_fallback_layers: tuple[str, ...] = () + if not requested: + return + modules = dict(self.named_modules()) + missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)] + local_layers = [ + name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes) + ] + self._local_w4a8_fallback_layers = tuple(local_layers) + invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)] + if invalid: + raise ValueError( + "w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; " + f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, " + "not checkpoint attn1.to_q/to_k/to_v names." + ) + + def _log_w4a8_fallback_load_summary(self, component_name: str) -> None: + """Log post-load transforms/cache readiness, not numerical acceptance.""" + requested = self._w4a8_fallback_layers + if not requested: + return + modules = dict(self.named_modules()) + local_layers = self._local_w4a8_fallback_layers + selected_layers = [ + name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False) + ] + bf16_overrides = [name for name in local_layers if name not in selected_layers] + processed = [ + name + for name in selected_layers + if getattr(modules[name], "_already_called_process_weights_after_loading", False) + ] + if not selected_layers: + processing_state = "not-required" + elif len(processed) == len(selected_layers): + processing_state = "ready" + else: + processing_state = "not-ready" + + def qualified(names: Iterable[str]) -> list[str]: + return [f"{component_name}.{name}" for name in names] + + logger.info( + "Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; " + "weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d", + component_name, + len(selected_layers), + qualified(selected_layers), + len(bf16_overrides), + qualified(bf16_overrides), + processing_state, + len(processed), + len(selected_layers), + len(requested) - len(local_layers), + ) @property def dtype(self) -> torch.dtype: @@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module): } params_dict = dict(self.named_parameters()) + single_scale_prefixes = { + name + for name, module in self.named_modules() + if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod) + } + + def reject_dualscale_tensor(name: str) -> None: + if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes: + raise ValueError( + f"Single-scale mxfp4 cannot load DualScale tensor {name}; " + "changing quant_method does not convert a checkpoint." + ) + loaded_params: set[str] = set() + loaded_qkv_shards: dict[str, set[str]] = {} + qkv_smooth_scales: dict[str, torch.Tensor] = {} for name, loaded_weight in weights: name = weight_name_remapping.get(name, name) @@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module): # Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall # through to the else branch and are loaded directly. for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in original_name: + if f"{weight_name}." not in original_name: continue - lookup_name = original_name.replace(weight_name, param_name) + lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1) + reject_dualscale_tensor(lookup_name) # Skip weights that belong to PP stages other than this one if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict: break param = params_dict[lookup_name] + if lookup_name.endswith(".mul_scale"): + previous_scale = qkv_smooth_scales.get(lookup_name) + if previous_scale is not None and not torch.equal(previous_scale, loaded_weight): + raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}") + if previous_scale is None: + qkv_smooth_scales[lookup_name] = loaded_weight.clone() weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) + loaded_params.add(original_name) + if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False): + # Shared input-channel Smooth is a complete tensor even + # when supplied under a single Q/K/V source name. + loaded_params.add(lookup_name) + else: + shards = loaded_qkv_shards.setdefault(lookup_name, set()) + shards.add(shard_id) + if shards == {"q", "k", "v"}: + loaded_params.add(lookup_name) break else: # diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight @@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module): if is_pp_missing_parameter(lookup_name, self): continue + reject_dualscale_tensor(lookup_name) if lookup_name not in params_dict: logger.warning(f"Skipping weight {original_name} -> {lookup_name}") continue @@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) - - loaded_params.add(original_name) - loaded_params.add(lookup_name) + loaded_params.update((original_name, lookup_name)) return loaded_params diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 8f15aaf50..c0975f651 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -446,6 +446,33 @@ def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> b return False +_MXFP4_SERIALIZED_FLAGS = { + "mxfp4": "is_checkpoint_mxfp4_serialized", + "mxfp4_dualscale": "is_checkpoint_serialized", +} + + +def _normalize_serialized_mxfp4_layer_policy(qc_method: str, qc_kwargs: dict[str, Any]) -> None: + """Canonicalize the BF16 layer policy owned by serialized MXFP4 metadata. + + A serialized checkpoint config is a complete storage contract: an omitted + layer list means no explicit BF16 layers. Method-only online metadata does + not carry that contract and must leave caller-provided layer policy intact. + """ + method = _normalize_quant_method_alias(qc_method) + if method is None: + return + serialized_flag = _MXFP4_SERIALIZED_FLAGS.get(method) + if serialized_flag is None or not qc_kwargs.get(serialized_flag): + return + + ignored_layers = qc_kwargs.get("ignored_layers") + if not ignored_layers: + ignored_layers = qc_kwargs.get("modules_to_not_convert") + qc_kwargs["ignored_layers"] = ignored_layers or [] + qc_kwargs.pop("modules_to_not_convert", None) + + def resolve_quant_config_from_disk( quant_config: QuantizationConfig | None, disk_qc: dict[str, Any] | str | None, @@ -460,22 +487,24 @@ def resolve_quant_config_from_disk( - quant_config is None: auto-detect from disk_qc (full build). - Methods mismatch: raise ValueError — prevents silent weight corruption. - Disk marks serialized but quant_config is online: rebuild from disk. - - ignored_layers differ: rebuild from disk (per-transformer BF16 routing). + - Serialized MXFP4 metadata owns the complete BF16 layer policy; omitted + ignored_layers means the checkpoint default (no explicit BF16 layers). + - Method-only metadata carries no storage policy and preserves caller + ignored_layers and runtime fallback policies. """ if disk_qc is None: return quant_config if isinstance(disk_qc, str): - if quant_config is None: - logger.info("Auto-detected quantization from config.json: method=%s", disk_qc) - return build_quant_config(disk_qc) - return quant_config + qc_method = disk_qc + qc_kwargs: dict[str, Any] = {} + else: + if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc: + return quant_config + qc_method = disk_qc["quant_method"] + qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"} - if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc: - return quant_config - - qc_method: str = disk_qc["quant_method"] - qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"} + _normalize_serialized_mxfp4_layer_policy(qc_method, qc_kwargs) if quant_config is None: logger.info( @@ -494,6 +523,26 @@ def resolve_quant_config_from_disk( "Pass a matching --quantization flag or omit it for auto-detection." ) + if isinstance(disk_qc, str): + return quant_config + + # Serialized checkpoint metadata owns storage and BF16 layer routing. The + # active config owns runtime step/layer policies, including explicit empty + # lists. Preserve them when each cascade transformer rebuilds from disk. + for policy in ("w4a8_fallback_steps", "w4a8_fallback_layers"): + if hasattr(quant_config, policy): + qc_kwargs[policy] = list(getattr(quant_config, policy)) + if hasattr(quant_config, "mxfp4_scale_alg"): + # Runtime activation policy is independent of how offline W4 was made. + qc_kwargs["mxfp4_scale_alg"] = quant_config.mxfp4_scale_alg + if hasattr(quant_config, "require_smooth_scale"): + disk_requires_smooth = qc_kwargs.get("require_smooth_scale", False) + if not isinstance(disk_requires_smooth, bool): + raise ValueError("require_smooth_scale must be a boolean") + # Either a calibrated checkpoint or an explicit caller can require + # Smooth. Per-expert rebuilds must never weaken that requirement. + qc_kwargs["require_smooth_scale"] = quant_config.require_smooth_scale or disk_requires_smooth + if _disk_marks_serialized(qc_kwargs, quant_config): logger.info( "config.json marks checkpoint as serialized; switching to offline %s mode.", @@ -501,6 +550,14 @@ def resolve_quant_config_from_disk( ) return build_quant_config(qc_method, **qc_kwargs) + if ( + "require_smooth_scale" in qc_kwargs + and hasattr(quant_config, "require_smooth_scale") + and qc_kwargs["require_smooth_scale"] != quant_config.require_smooth_scale + ): + logger.info("config.json Smooth requirement differs from active config; rebuilding quant_config.") + return build_quant_config(qc_method, **qc_kwargs) + # AutoRound MXFP checkpoints use data_type="mx_fp" instead of # is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is # selected according to the checkpoint's bit width. diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py index ac93eaa00..43f3efa9d 100644 --- a/vllm_omni/quantization/mxfp4_config.py +++ b/vllm_omni/quantization/mxfp4_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers. Architecture mirrors mxfp8_config.py: @@ -41,7 +41,7 @@ Dual-scale (W4A4_MXFP4_DUALSCALE): and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast. Checkpoint tensor shapes for dual-scale: - weight : (N, K) float8_e4m3fn – FP4 packed + weight : (N, K) float8_e4m3fn – numeric FP4 values (packed after loading) weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits) weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert) mul_scale : (K,) float32 – per-input-channel activation pre-scale @@ -51,9 +51,9 @@ Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/q from __future__ import annotations -import re -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast +import regex as re import torch from torch.nn import Module from vllm.logger import init_logger @@ -83,9 +83,54 @@ from vllm_omni.quantization.mxfp8_config import ( if TYPE_CHECKING: from vllm.model_executor.models.utils import WeightsMapper + logger = init_logger(__name__) +def _validate_w4a8_fallback_steps(steps: list[int] | None) -> list[int]: + if steps is None: + return [] + if not isinstance(steps, list) or any(type(step) is not int or step < 0 for step in steps): + raise ValueError("w4a8_fallback_steps must be a list of non-negative integer denoise step indices.") + return sorted(set(steps)) + + +def _validate_w4a8_fallback_layers(layers: list[str] | None) -> list[str]: + """Names are exact runtime Linear paths, not checkpoint aliases or patterns.""" + if layers is None: + return [] + if not isinstance(layers, list) or any( + not isinstance(name, str) or re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*(?:\.[A-Za-z_0-9]+)*", name) is None + for name in layers + ): + raise ValueError("w4a8_fallback_layers must be a list of non-empty exact runtime Linear paths.") + return sorted(set(layers)) + + +def _is_w4a8_fallback_step(steps: list[int]) -> bool: + """Read the request's zero-based denoise index; never keep state on a Linear.""" + if not steps: + return False + + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + step = get_forward_context().denoise_step_idx if is_forward_context_available() else None + if step is None: + raise RuntimeError("W4A8 step fallback requires a pipeline that publishes its denoise step in ForwardContext.") + return step in steps + + +def _validate_smooth_scale(scale: torch.Tensor) -> None: + """Reject malformed calibration before loading or activation dtype conversion.""" + if scale.ndim != 1 or not scale.is_floating_point(): + raise ValueError("MXFP4 mul_scale must be a one-dimensional floating-point Smooth tensor.") + if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()): + raise ValueError("MXFP4 mul_scale must contain finite, strictly positive Smooth values.") + + # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @@ -105,10 +150,24 @@ class DiffusionMXFP4Config(QuantizationConfig): self, is_checkpoint_mxfp4_serialized: bool = False, ignored_layers: list[str] | None = None, + w4a8_fallback_steps: list[int] | None = None, + require_smooth_scale: bool = False, + w4a8_fallback_layers: list[str] | None = None, + mxfp4_scale_alg: int = 0, ) -> None: super().__init__() self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized self.ignored_layers = ignored_layers or [] + self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps) + self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers) + if not isinstance(require_smooth_scale, bool): + raise ValueError("require_smooth_scale must be a boolean") + if require_smooth_scale and not self.is_checkpoint_mxfp4_serialized: + raise ValueError("require_smooth_scale requires an offline MXFP4 checkpoint") + self.require_smooth_scale = require_smooth_scale + if type(mxfp4_scale_alg) is not int or mxfp4_scale_alg not in (0, 2): + raise ValueError("mxfp4_scale_alg must be 0 (OCP MX) or 2 (UOS with dst_type_max=7.25).") + self.mxfp4_scale_alg = mxfp4_scale_alg @classmethod def get_name(cls) -> QuantizationMethods: @@ -139,6 +198,10 @@ class DiffusionMXFP4Config(QuantizationConfig): return cls( is_checkpoint_mxfp4_serialized=is_serialized, ignored_layers=ignored_layers, + w4a8_fallback_steps=config.get("w4a8_fallback_steps"), + w4a8_fallback_layers=config.get("w4a8_fallback_layers"), + require_smooth_scale=config.get("require_smooth_scale", False), + mxfp4_scale_alg=config.get("mxfp4_scale_alg", 0), ) def get_quant_method( @@ -155,9 +218,15 @@ class DiffusionMXFP4Config(QuantizationConfig): return UnquantizedLinearMethod() if current_omni_platform.is_npu(): if self.is_checkpoint_mxfp4_serialized: - return NPUMxfp4LinearMethod(self) - return NPUMxfp4OnlineLinearMethod(self) + return NPUMxfp4LinearMethod(self, prefix=prefix) + return NPUMxfp4OnlineLinearMethod(self, prefix=prefix) if current_omni_platform.is_rocm(): + if self.mxfp4_scale_alg != 0: + raise NotImplementedError("MXFP4 UOS quantization is currently only supported on NPU (Ascend).") + if self.w4a8_fallback_steps or self.w4a8_fallback_layers: + raise NotImplementedError( + "MXFP4 W4A8 step/layer fallback is currently only supported on NPU (Ascend)." + ) gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName if "gfx950" not in gcn_arch: raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}") @@ -176,6 +245,56 @@ class DiffusionMXFP4Config(QuantizationConfig): # --------------------------------------------------------------------------- +def _npu_quantize_mxfp4(x: torch.Tensor, scale_alg: int) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight or activation using the same explicit FP4 algorithm.""" + import torch_npu + + # Universal Optimal Scaling (Qmax=7.25), MXAttention Section 4.1: + # https://arxiv.org/abs/2607.24377 + scale_kwargs = {"dst_type_max": 7.25} if scale_alg == 2 else {} + return torch_npu.npu_dynamic_mx_quant( + x, + dst_type=torch_npu.float4_e2m1fn_x2, + axis=-1, + block_size=32, + round_mode="rint", + scale_alg=scale_alg, + **scale_kwargs, + ) + + +def _npu_w4a8_matmul( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, +) -> torch.Tensor: + """MXFP8 activations with prepared single-scale MXFP4 weights (N, K).""" + import torch_npu + + x_q, x_scale = torch_npu.npu_dynamic_mx_quant( + x, dst_type=torch_npu.float8_e4m3fn, axis=-1, block_size=32, round_mode="rint", scale_alg=0 + ) + # Use the BF16 output/bias pair supported by the mixed FP8/FP4 kernel + # in the target A5 runtime. Restore the caller's dtype afterwards. + if bias is not None: + bias = bias.to(torch.bfloat16).reshape(1, -1) + output = torch_npu.npu_quant_matmul( + x_q, + weight.transpose(0, 1), + weight_scale.transpose(0, 1), + scale_dtype=torch_npu.float8_e8m0fnu, + x2_dtype=torch_npu.float4_e2m1fn_x2, + pertoken_scale=x_scale, + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, + bias=bias, + output_dtype=torch.bfloat16, + group_sizes=[0, 0, 32], + ) + return output.to(output_dtype) + + class NPUMxfp4LinearMethod(MXFPLinearMethodBase): """NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints. @@ -187,7 +306,8 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared. """ - def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + def __init__(self, quant_config: DiffusionMXFP4Config, *, prefix: str = "") -> None: + self.is_w4a8_fallback_layer = prefix in quant_config.w4a8_fallback_layers self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() @@ -210,11 +330,14 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): layer.orig_dtype = params_dtype layer.weight_block_size = None - # BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights. layer.register_parameter( "weight", ModelWeightParameter( - data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=params_dtype, + ), input_dim=1, output_dim=0, weight_loader=weight_loader, @@ -234,12 +357,42 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): ), ) + # Older single-scale checkpoints have no Smooth tensor. Keep their + # identity transform while allowing calibrated checkpoints to load it. + def smooth_weight_loader( + param: ModelWeightParameter, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any + ) -> None: + _validate_smooth_scale(loaded_weight) + if not callable(weight_loader): + raise ValueError("Loading MXFP4 Smooth requires a callable weight_loader") + weight_loader(param, loaded_weight, *args, **kwargs) + + layer.register_parameter( + "mul_scale", + ModelWeightParameter( + data=torch.ones(input_size_per_partition, dtype=torch.float32), + input_dim=0, + output_dim=None, + weight_loader=smooth_weight_loader, + ), + ) + setattr(layer.mul_scale, "ignore_warning", True) + setattr(layer.mul_scale, "is_checkpoint_optional", not self.quant_config.require_smooth_scale) + if self.quant_config.is_checkpoint_mxfp4_serialized: + setattr(layer.weight_scale, "is_checkpoint_required", True) + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return import torch_npu + smooth_param = cast(torch.Tensor, layer.mul_scale) + _validate_smooth_scale(smooth_param) + smooth = smooth_param.data.to(dtype=cast(torch.dtype, layer.orig_dtype)).contiguous() + # Finite FP32 values can overflow or underflow in the activation dtype. + _validate_smooth_scale(smooth) + # NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose. w = layer.weight if w.dtype != torch_npu.float4_e2m1fn_x2: @@ -258,15 +411,31 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): replace_parameter(layer, "weight", w) replace_parameter(layer, "weight_scale", s) + # Keep loader metadata while converting the calibration tensor once. + layer.mul_scale.data = smooth.to(device=cast(torch.Tensor, w).device) layer._already_called_process_weights_after_loading = True # --- NPU MXFP4 ops — shared with online path via inheritance --- - def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - import torch_npu + def _apply_inner( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + # Online quantization has no calibrated Smooth tensor. For offline + # weights the same pre-scale must precede both A4 and A8 quantization. + if (mul_scale := getattr(layer, "mul_scale", None)) is not None: + x = x * mul_scale.to(x.dtype) + if self.is_w4a8_fallback_layer or _is_w4a8_fallback_step(self.quant_config.w4a8_fallback_steps): + return _npu_w4a8_matmul( + x, cast(torch.Tensor, layer.weight), cast(torch.Tensor, layer.weight_scale), bias, ori_dtype + ) + return super()._apply_inner(layer, x, bias, ori_dtype) - # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2. - return torch_npu.npu_dynamic_mx_quant(x) + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return _npu_quantize_mxfp4(x, self.quant_config.mxfp4_scale_alg) def _quant_matmul( self, @@ -318,8 +487,6 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): if getattr(layer, "_already_called_process_weights_after_loading", False): return - import torch_npu - if layer.weight.device == torch.device("meta"): weight = ModelWeightParameter( data=torch.empty_like(layer.weight, device=layer._load_device), @@ -331,8 +498,10 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): layer.register_parameter("weight", weight) initialize_single_dummy_weight(layer.weight) - # NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2. - weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight) + # Prepare once; offline checkpoints keep their existing weight algorithm. + weight_fp4, weight_scale_raw = _npu_quantize_mxfp4( + cast(torch.Tensor, layer.weight), self.quant_config.mxfp4_scale_alg + ) # Weight stays (N, K) — no pre-transpose for FP4 packed format. # Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul. @@ -497,7 +666,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): """NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints. Checkpoint tensors and their canonical post-load shapes: - weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte) + weight : (N, K) float8_e4m3fn – numeric FP4 values; packed into NZ after loading weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2) weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N) mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration) @@ -509,7 +678,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): Reference: MindIE-SD W4A4MXFP4DualQuantLinear. """ - def __init__(self, quant_config: Any) -> None: + def __init__(self, quant_config: DiffusionMXFP4DualScaleMixedConfig) -> None: self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() @@ -532,7 +701,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): layer.orig_dtype = params_dtype layer.weight_block_size = None - # FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16 + # Checkpoint FP4 values are stored numerically in FP8; BF16 preserves them exactly. layer.register_parameter( "weight", ModelWeightParameter( @@ -587,6 +756,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): ), ) setattr(layer.mul_scale, "ignore_warning", True) + setattr(layer.weight_scale, "is_checkpoint_required", True) def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): @@ -594,7 +764,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): import torch_npu - # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29). + # Numeric FP4 values → packed float4_e2m1fn_x2 → NZ hardware format (format ID 29). # NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices. w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2) w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8) @@ -631,7 +801,10 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): """ if ori_dtype not in (torch.bfloat16, torch.float16): x = x.to(torch.bfloat16) - x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale) + # The dual-level activation kernel requires x and Smooth to have the + # same dtype. Prepared calibration scales are BF16, including online. + smooth_scale = cast(torch.Tensor, layer.mul_scale).to(x.dtype) + x_q, l0_scale, l1_scale = self._quantize_activation(x, smooth_scale) return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype) def _quantize_activation( # type: ignore[override] @@ -785,11 +958,20 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): is_checkpoint_serialized: bool = False, ignored_layers: list[str] | None = None, num_bf16_fallback_layers: int = 5, + w4a8_fallback_steps: list[int] | None = None, + w4a8_fallback_layers: list[str] | None = None, ) -> None: super().__init__() self.is_checkpoint_serialized = is_checkpoint_serialized self.ignored_layers = ignored_layers or [] self.num_bf16_fallback_layers = num_bf16_fallback_layers + self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps) + self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers) + if self.w4a8_fallback_steps or self.w4a8_fallback_layers: + raise ValueError( + "mxfp4_dualscale does not support W4A8 fallback. Use mxfp4 with single-scale weights; " + "changing a DualScale checkpoint's method name does not convert its storage format." + ) @classmethod def get_name(cls) -> QuantizationMethods: @@ -822,6 +1004,8 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): is_checkpoint_serialized=is_serialized, ignored_layers=ignored_layers, num_bf16_fallback_layers=num_bf16_fallback_layers, + w4a8_fallback_steps=config.get("w4a8_fallback_steps"), + w4a8_fallback_layers=config.get("w4a8_fallback_layers"), ) def get_quant_method( From f1a6e7ce2a8fceb1417ec86b0b49a6386601ac65 Mon Sep 17 00:00:00 2001 From: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:47:38 +0900 Subject: [PATCH 20/99] [Core] Retain structured configs through runtime startup (#6849) Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> --- docs/design/architecture_overview.md | 12 +- .../test_factory_integration.py | 31 ++- tests/config/test_forced_aligner_injection.py | 36 ++++ tests/config/test_omni_config.py | 25 +++ .../test_async_omni_qwen_image_generate.py | 10 + .../offline_inference/test_hunyuanimage3.py | 162 +++++++++++++--- .../e2e/offline_inference/test_qwen3_omni.py | 85 +++++++++ tests/e2e/offline_inference/test_wan22_t2v.py | 26 +++ .../test_async_omni_engine_abort_ack.py | 5 +- .../test_async_omni_engine_stage_init.py | 178 ++++++++++++++++++ tests/engine/test_parallel_stage_init.py | 17 ++ tests/engine/test_single_stage_mode.py | 21 ++- tests/engine/test_stage_device_layout.py | 37 ++++ tests/engine/test_stage_engine_args.py | 16 ++ .../openai_api/test_api_server_guards.py | 7 + .../openai_api/test_audex_serving_guards.py | 45 +++++ .../openai_api/test_openpi_serving.py | 30 +++ ...serving_chat_minicpmo45_reference_audio.py | 33 +++- ...ing_chat_minicpmo45_stream_output_kinds.py | 22 +++ .../openai_api/test_serving_speech.py | 33 ++++ .../openai_api/test_tts_adapter.py | 26 +++ .../openai_api/test_video_server.py | 55 ++++++ .../test_async_omni_diffusion_config.py | 136 ++++++++++++- tests/entrypoints/test_omni_entrypoints.py | 62 ++++-- .../test_resolve_dreamzero_config.py | 9 +- tests/entrypoints/test_serve.py | 93 +++++++++ tests/entrypoints/test_utils.py | 174 ++++++++++++----- vllm_omni/config/composable_parallel/apply.py | 38 +++- vllm_omni/config/config_factory.py | 92 ++++++++- vllm_omni/config/omni_config.py | 136 ++++++++++++- vllm_omni/config/resolver.py | 51 ++--- vllm_omni/engine/omni_engine_base.py | 9 +- vllm_omni/engine/stage_admission.py | 7 +- vllm_omni/engine/stage_engine_startup.py | 17 +- vllm_omni/engine/stage_init_utils.py | 116 ++++++++---- vllm_omni/engine/stage_runtime.py | 93 +++++---- vllm_omni/entrypoints/cli/serve.py | 27 +-- vllm_omni/entrypoints/openai/serving_chat.py | 26 ++- .../entrypoints/openai/serving_speech.py | 46 +++-- vllm_omni/entrypoints/openai/serving_video.py | 50 ++--- .../entrypoints/openai/tts_adapters/audex.py | 6 +- .../entrypoints/openai/tts_adapters/base.py | 15 ++ .../openai/tts_adapters/higgs_audio_v2.py | 19 +- .../openai/tts_adapters/higgs_audio_v3.py | 9 +- .../openai/video/generation/helpers.py | 7 +- vllm_omni/entrypoints/openpi/serving.py | 8 + vllm_omni/entrypoints/serve/profile/utils.py | 16 +- vllm_omni/entrypoints/stage_utils.py | 6 + vllm_omni/entrypoints/utils.py | 7 + 49 files changed, 1856 insertions(+), 331 deletions(-) diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 862516831..8b7168824 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -323,7 +323,7 @@ inputs, fields, and ownership rules are described below. flowchart TB layer1["Layer 1 · Authoring inputs
PipelineConfig + DeployConfig"] layer2["Layer 2 · Production resolution boundary
resolve_omni_config()
(StageConfigFactory + VllmOmniConfig internally)"] - layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + temporary OmegaConf stage bridge"] + layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + typed stage configs"] layer4["Layer 4 · Runtime launch planning
StageRuntime"] layer5["Layer 5 · Engine materialization
VllmConfig / OmniDiffusionConfig"] @@ -335,10 +335,10 @@ The production resolution boundary is `resolve_omni_config()` in It delegates typed construction to `StageConfigFactory.create_from_model()` and `VllmOmniConfig.from_pipeline_config()`, then returns an `OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup. -Until `StageRuntime` consumes typed stage configs directly, this envelope carries -the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as -a temporary runtime bridge; both views describe the same resolved topology, -including injected stages. +This envelope carries the effective `PipelineConfig` alongside typed +`stage_configs`, which `StageRuntime` consumes directly. Both views describe the +same resolved topology, including injected stages. Backend arguments are +projected from these typed stages when each engine is initialized. The legacy `stage_args` YAML path has been removed. Model topology now resolves through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`. @@ -360,7 +360,7 @@ The important ownership rules are: 3. CLI and Python overrides are applied at the resolution boundary, with per-stage overrides taking precedence over global values where supported. 4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config` - and temporary `stage_configs` compatibility view must describe the same + and typed `stage_configs` must describe the same topology. 5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan` is runtime-private state, not a user configuration object. diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py index 468b3db2c..d49d32042 100644 --- a/tests/config/composable_parallel/test_factory_integration.py +++ b/tests/config/composable_parallel/test_factory_integration.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Integration tests: strategy overlay on the real merge_pipeline_deploy seam.""" from __future__ import annotations @@ -20,6 +20,7 @@ from vllm_omni.config.composable_parallel import ( apply_strategy_specs, ) from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy @@ -137,3 +138,31 @@ def test_cli_overrides_strategy_with_warning(): assert _resolved(stages, "talker").runtime.num_replicas == 3 # ...and the override was warned about, naming the conflicting field. assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages) + + +def test_typed_cli_overrides_strategy(): + pipeline = OMNI_PIPELINES["qwen2_5_omni"] + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=load_deploy_config(_DEPLOY), + cli_overrides={"stage_1_num_replicas": 3}, + strategy_specs={"talker": [_stage_replica(2, "round_robin")]}, + ) + + talker = next(stage for stage in config.stage_configs if stage.model_stage == "talker") + assert talker.runtime_config.num_replicas == 3 + assert config.orchestrator_config.omni_lb_policy == "round-robin" + assert config.strategy_omni_lb_policy == "round-robin" + + +def test_typed_tp_only_strategy_has_no_derived_lb_policy(): + pipeline = OMNI_PIPELINES["qwen2_5_omni"] + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=load_deploy_config(_DEPLOY), + cli_overrides={"omni_lb_policy": "round-robin"}, + strategy_specs={"thinker": [_tp(1)]}, + ) + + assert config.orchestrator_config.omni_lb_policy == "round-robin" + assert config.strategy_omni_lb_policy is None diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py index 65c315229..2037b22f1 100644 --- a/tests/config/test_forced_aligner_injection.py +++ b/tests/config/test_forced_aligner_injection.py @@ -2,13 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest +from vllm.pooling_params import PoolingParams import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry) +from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY from vllm_omni.config.stage_config import ( DeployConfig, StageExecutionType, ) +from vllm_omni.engine.stage_init_utils import ( + build_engine_args_dict_from_omni_stage_config, + extract_stage_metadata_from_omni_stage_config, +) from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( POOLING_OUTPUT_DECODER_PATH, TIMESTAMPS_MODALITY, @@ -94,3 +100,33 @@ def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_ch assert stages[2].engine_args.runner == "pooling" assert len(pipeline.stages) == 2 assert deploy.stages == [] + + +def test_injected_aligner_survives_typed_config_and_runtime_projections(): + pipeline, deploy = inject_forced_aligner_stage( + _PIPELINE_REGISTRY["qwen3_tts"], + DeployConfig(), + {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}, + ) + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=deploy, + cli_overrides={"model": "/models/Qwen3-TTS"}, + ) + stage = config.stage_configs[-1] + + assert stage.pooling_config.runner == "pooling" + assert stage.pooling_config.pooling_output_decoder == POOLING_OUTPUT_DECODER_PATH + assert isinstance(stage.pooling_config.default_pooling_params, PoolingParams) + assert stage.pooling_config.default_pooling_params.task == deploy.stages[-1].default_pooling_params["task"] + + metadata = extract_stage_metadata_from_omni_stage_config(stage) + assert isinstance(metadata.default_sampling_params, PoolingParams) + assert metadata.default_sampling_params.task == deploy.stages[-1].default_pooling_params["task"] + + engine_args = build_engine_args_dict_from_omni_stage_config( + stage, + stage.model_config.model or "/models/Qwen3-ForcedAligner-0.6B", + ) + assert engine_args["runner"] == "pooling" + assert engine_args["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1ed62f7f1..97f9a5370 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -11,6 +11,7 @@ from pathlib import Path import msgspec import pytest +import torch from pydantic import ValidationError from pydantic.fields import FieldInfo from transformers import Qwen3OmniMoeConfig @@ -123,6 +124,18 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) - assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"] +def test_nested_stage_override_deep_merges_structured_model_config() -> None: + config = _from_pipeline_key( + "cosmos3_policy", + deploy_config_path=get_deploy_config_path("cosmos3_policy_droid.yaml"), + cli_overrides={"stage_0_model_config": {"guardrails": False}}, + ) + + model_config = config.stage_by_id(0).diffusion_config.model_config + assert model_config["guardrails"] is False + assert model_config["policy_server_config"]["action_space"] == "joint_position" + + @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str): pipeline = _resolve_pipeline_or_skip(model_type) @@ -654,6 +667,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): "load_config", "cache_config", "scheduler_config", + "pooling_config", "connector_config", "runtime_config", "parallel_config", @@ -889,6 +903,17 @@ def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_saf msgspec.msgpack.encode(payload) +def test_structured_diffusion_torch_dtype_is_msgpack_transport_safe(): + stage_config = _from_pipeline_key( + "hunyuan_image3_dit", + cli_overrides={"dtype": torch.bfloat16}, + ).stage_by_id(0) + + serialized = _serialize_stage_config(stage_config) + assert serialized["diffusion_config"]["dtype"] == "bfloat16" + msgspec.msgpack.encode(serialized) + + def test_diffusion_parallel_config_fields_cover_legacy_surface(): from vllm_omni.diffusion.data import DiffusionParallelConfig diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index 56a960121..4bea6563b 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -21,6 +21,7 @@ import torch.distributed as dist from transformers import AutoTokenizer from tests.helpers.mark import hardware_test +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -136,6 +137,15 @@ async def _generate_once( def _assert_live_uniproc_worker(engine: AsyncOmni) -> None: """Default single-GPU path must use the in-process executor, not mp IPC.""" + stage_configs = engine.engine.stage_configs + assert len(stage_configs) == 1 + stage_config = stage_configs[0] + assert isinstance(stage_config, VllmOmniDiffusionStageConfig) + assert stage_config.model_config.model == MODEL + assert stage_config.model_config.enforce_eager is True + assert stage_config.scheduler_config.max_num_seqs == 1 + assert stage_config.diffusion_config.custom_pipeline_args == {"pipeline_class": CUSTOM_PIPELINE_CLASS} + assert stage_config.diffusion_config.worker_extension_cls == WORKER_EXTENSION_CLASS clients = engine.engine.stage_clients assert len(clients) == 1 client = clients[0] diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index 8655bd63f..21bfe9b39 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -1,5 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # ruff: noqa: E501 +import copy import importlib.util +import os import sys from collections.abc import Generator from pathlib import Path @@ -10,20 +15,37 @@ import torch.nn.functional as F from PIL import Image from transformers import CLIPModel, CLIPProcessor +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni import Omni +from vllm_omni.config.omni_config import ( + VllmOmniARStageConfig, + VllmOmniDiffusionStageConfig, +) +from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name from vllm_omni.platforms import current_omni_platform PROMPT = "A brown and white dog is running on the grass" -MODEL_NAME = "tencent/HunyuanImage-3.0" +MODEL_NAME = os.environ.get("HUNYUAN_IMAGE3_MODEL", "tencent/HunyuanImage-3.0-Instruct") LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" -DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") +_PRODUCTION_DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") +# Keep the production Mooncake connector, but use its TCP transport for this +# single-host test. The shared server's RDMA device is not usable reliably; +# TCP still exercises the Mooncake AR→DiT data plane without changing deploy +# semantics in production. +DEPLOY_CONFIG_PATH = modify_stage_config( + _PRODUCTION_DEPLOY_CONFIG_PATH, + updates={"connectors.rdma_connector.extra.protocol": "tcp"}, +) -pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] +pytestmark = [ + pytest.mark.advanced_model, + pytest.mark.diffusion, +] # System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified # Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass"). @@ -285,25 +307,97 @@ def omni() -> Generator[Omni, None, None]: yield runner.omni +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") +@hardware_test(res={"cuda": "H100"}, num_cards=4) +def test_structured_mixed_pipeline_config_reaches_runtime(omni: Omni) -> None: + """Mixed-pipeline deploy settings reach both terminal engine configs.""" + engine = omni.engine + + # Topology, connector selection, and placement remain Omni-owned. + resolved_stages = engine.stage_configs + assert len(resolved_stages) == 2 + ar_stage, diffusion_stage = resolved_stages + assert isinstance(ar_stage, VllmOmniARStageConfig) + assert ar_stage.stage_id == 0 + assert ar_stage.model_stage == "AR" + assert ar_stage.final_output_type == "text" + assert ar_stage.runtime_config.devices == "0,1" + assert ar_stage.connector_config.output_connectors == {"to_stage_1": "rdma_connector"} + assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig) + assert diffusion_stage.stage_id == 1 + assert diffusion_stage.model_stage == "dit" + assert diffusion_stage.final_output_type == "image" + assert diffusion_stage.input_sources == [0] + assert diffusion_stage.runtime_config.devices == "2,3" + assert diffusion_stage.connector_config.input_connectors == {"from_stage_0": "rdma_connector"} + + # Stage 0 is materialized as the VllmConfig actually passed to its engine. + ar_vllm_config, diffusion_vllm_config = engine.stage_vllm_configs + assert ar_vllm_config is not None + assert diffusion_vllm_config is None + assert ar_vllm_config.model_config.enforce_eager is True + assert ar_vllm_config.model_config.trust_remote_code is True + assert ar_vllm_config.parallel_config.tensor_parallel_size == 2 + assert ar_vllm_config.scheduler_config.max_num_seqs == 1 + assert ar_vllm_config.scheduler_config.max_num_batched_tokens == 32768 + assert ar_vllm_config.cache_config.gpu_memory_utilization == 0.95 + ar_omni_config = ar_vllm_config.model_config.omni_config + assert ar_omni_config.omni_kv_config["need_send_cache"] is True + + expected_ar_sampling = { + "temperature": 0.0, + "top_p": 1, + "top_k": -1, + "max_tokens": 8192, + "detokenize": True, + "skip_special_tokens": False, + "include_stop_str_in_output": True, + } + ar_sampling, diffusion_sampling = engine.default_sampling_params_list + assert all(getattr(ar_sampling, name) == value for name, value in expected_ar_sampling.items()) + assert diffusion_sampling.num_inference_steps == 50 + assert diffusion_sampling.guidance_scale == 0 + + # Stage 1 is a diffusion engine, so its terminal owner is the launched + # OmniDiffusionConfig rather than a VllmConfig. + od_config = omni.get_diffusion_od_config() + assert od_config is not None + assert od_config.model == MODEL_NAME + assert od_config.enforce_eager is True + assert od_config.distributed_executor_backend == "mp" + assert od_config.parallel_config.tensor_parallel_size == 2 + assert od_config.parallel_config.enable_expert_parallel is True + assert od_config.max_num_seqs == 1 + diffusion_kv = od_config.omni_kv_config + assert diffusion_kv["need_recv_cache"] is True + assert diffusion_kv["enable_kv_async_prefetch"] is True + assert diffusion_kv["kv_prefetch_min_free_mem_ratio"] == 0.2 + + def _extract_generated_image(outputs: list[object]) -> Image.Image: if not outputs: raise AssertionError("No outputs were returned from Omni.generate()") - first_output = outputs[0] - if images := getattr(first_output, "images", None): - return images[0] - - request_output = first_output - if request_output is not None and (images := getattr(request_output, "images", None)): - return images[0] + for output in outputs: + if images := getattr(output, "images", None): + return images[0] + multimodal_output = getattr(output, "multimodal_output", None) + if isinstance(multimodal_output, dict) and (images := multimodal_output.get("images")): + return images[0] raise AssertionError("No generated image found in Omni output") def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor: - inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt") + # The full CLIP forward API is stable across supported Transformers + # versions, but it expects text inputs as well as pixels. An empty text + # prompt is sufficient when only the projected image embedding is needed. + inputs = clip_processor(text=[""], images=[image.convert("RGB")], return_tensors="pt", padding=True) with torch.inference_mode(): - features = clip_model.get_image_features(**inputs) + # Use the public CLIP forward output, matching the shared CLIPScorer + # helper. ``get_image_features`` changed from Tensor to + # BaseModelOutputWithPooling in Transformers 5.x. + features = clip_model(**inputs).image_embeds features = F.normalize(features, p=2, dim=-1) return features.squeeze(0) @@ -320,16 +414,36 @@ def compare_semantic( def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: - generator_device = current_omni_platform.device_type or "cuda" - sampling_params = OmniDiffusionSamplingParams( - seed=1234, - generator=torch.Generator(device=generator_device).manual_seed(1234), - num_outputs_per_prompt=1, - ) - if use_system_prompt is not None: - sampling_params.extra_args = {"use_system_prompt": use_system_prompt} + from transformers import AutoTokenizer - outputs = omni.generate({"prompt": PROMPT}, sampling_params) + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + prompt_tokens = build_prompt_tokens( + PROMPT, + tokenizer, + task="t2i", + bot_task="think", + sys_type=use_system_prompt or "None", + ) + stop_token_ids = resolve_stop_token_ids(task="t2i", bot_task="think", tokenizer=tokenizer, image_size="1024x1024") + params_list = copy.deepcopy(omni.default_sampling_params_list) + generator_device = current_omni_platform.device_type or "cuda" + for params in params_list: + if isinstance(params, OmniDiffusionSamplingParams): + params.seed = 1234 + params.generator = torch.Generator(device=generator_device).manual_seed(1234) + params.num_outputs_per_prompt = 1 + elif hasattr(params, "stop_token_ids"): + params.stop_token_ids = stop_token_ids + + prompt = { + "prompt": PROMPT, + "prompt_token_ids": prompt_tokens.token_ids, + "use_system_prompt": prompt_tokens.system_prompt_type, + "modalities": ["image"], + "height": 1024, + "width": 1024, + } + outputs = omni.generate(prompt, params_list) return _extract_generated_image(outputs) @@ -342,7 +456,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: MIN_SEMANTIC_SIMILARITY = 0.5 -@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES) def test_system_prompt_scores( omni: Omni, @@ -378,7 +492,7 @@ def _load_text_to_image_module(): return module -@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") def test_shared_script_ar_path_reaches_generation( omni: Omni, clip_bundle: tuple[CLIPModel, CLIPProcessor], diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 32c345194..f5937265c 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -17,6 +17,10 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_video from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.config.omni_config import ( + VllmOmniARStageConfig, + VllmOmniGenerationStageConfig, +) from vllm_omni.platforms import current_omni_platform models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -25,6 +29,7 @@ thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"] # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the # platforms: section. Only CUDA needs an extra enforce_eager tweak. _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") +_PRODUCTION_DEPLOY = get_deploy_config_path("qwen3_omni_moe.yaml") def get_cuda_graph_config(): @@ -46,6 +51,7 @@ else: # Create parameter combinations for model and stage config test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +production_test_params = [(model, _PRODUCTION_DEPLOY) for model in models] # we can use the same config for a model that only has thinker (i.e., does not # enable audio output) because the resolver should figure out that it doesn't # need the full pipeline based on the HF config. @@ -59,6 +65,85 @@ def get_question(prompt_type="video"): return prompts.get(prompt_type, prompts["video"]) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", production_test_params, indirect=True) +def test_structured_multistage_config_reaches_runtime(omni_runner, offline_client) -> None: + """Deploy settings reach the materialized configs used by live stages.""" + engine = omni_runner.omni.engine + + # The resolved typed stages remain the Omni-owned source for topology and + # placement, which are not fields of the terminal vLLM config. + resolved_stages = engine.stage_configs + assert len(resolved_stages) == 3 + thinker, talker, code2wav = resolved_stages + assert isinstance(thinker, VllmOmniARStageConfig) + assert isinstance(talker, VllmOmniARStageConfig) + assert isinstance(code2wav, VllmOmniGenerationStageConfig) + assert [stage.stage_id for stage in resolved_stages] == [0, 1, 2] + assert [stage.model_stage for stage in resolved_stages] == ["thinker", "talker", "code2wav"] + assert code2wav.final_output_type == "audio" + assert thinker.runtime_config.devices == "0" + assert talker.runtime_config.devices == code2wav.runtime_config.devices == "1" + assert talker.connector_config.input_connectors == {"from_stage_0": "connector_of_shared_memory"} + assert code2wav.connector_config.input_connectors == {"from_stage_1": "connector_of_shared_memory"} + + # Engine-owned deploy settings must survive the startup projection into + # the actual VllmConfig objects retained by the stage pools. + assert len(engine.stage_vllm_configs) == 3 + thinker_vllm, talker_vllm, code2wav_vllm = engine.stage_vllm_configs + assert thinker_vllm is not None + assert talker_vllm is not None + assert code2wav_vllm is not None + assert thinker_vllm.scheduler_config.max_num_seqs == 64 + assert talker_vllm.scheduler_config.max_num_seqs == 64 + assert code2wav_vllm.scheduler_config.max_num_seqs == 64 + assert thinker_vllm.scheduler_config.max_num_batched_tokens == 32768 + assert talker_vllm.scheduler_config.max_num_batched_tokens == 32768 + assert code2wav_vllm.scheduler_config.max_num_batched_tokens == 65536 + assert thinker_vllm.cache_config.gpu_memory_utilization == 0.9 + assert talker_vllm.cache_config.gpu_memory_utilization == 0.6 + assert code2wav_vllm.cache_config.gpu_memory_utilization == 0.1 + assert all(config.cache_config.enable_prefix_caching is False for config in engine.stage_vllm_configs) + assert all(config.model_config.trust_remote_code is True for config in engine.stage_vllm_configs) + assert code2wav_vllm.scheduler_config.enable_chunked_prefill is False + assert code2wav_vllm.scheduler_config.async_scheduling is False + + # Sampling defaults are consumed from StageClient metadata, rather than + # from either the resolver output or VllmConfig. + expected_sampling = ( + {"temperature": 0.0, "max_tokens": 2048}, + {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "repetition_penalty": 1.05}, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "repetition_penalty": 1.1, + }, + ) + assert len(engine.default_sampling_params_list) == len(expected_sampling) + for runtime_params, expected in zip(engine.default_sampling_params_list, expected_sampling, strict=True): + # SamplingParams normalization may add backend defaults (for example + # ``detokenize=True``). Verify deploy-owned values without rejecting + # those non-lossy normalized fields. + assert all(getattr(runtime_params, name) == value for name, value in expected.items()) + + if current_omni_platform.is_cuda(): + assert thinker_vllm.model_config.enforce_eager is False + assert talker_vllm.model_config.enforce_eager is False + + # Exercise the materialized stages and their connector path with a real + # request; the assertions above are not merely resolver-shape checks. + offline_client.send_omni_request( + { + "prompts": "Answer with one short sentence: what is the capital of China?", + "modalities": ["text"], + } + ) + + @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py index 805fbd1fe..07ac5c8a6 100644 --- a/tests/e2e/offline_inference/test_wan22_t2v.py +++ b/tests/e2e/offline_inference/test_wan22_t2v.py @@ -5,6 +5,7 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OfflineOmniClient +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" @@ -12,6 +13,31 @@ PROMPT = "Two anthropomorphic cats in boxing gear on a spotlighted stage." NEGATIVE_PROMPT = "low quality, blurry, watermark, text" +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1}) +@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) +def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None: + """The default WAN resolver launches the typed diffusion config itself.""" + omni = omni_runner.omni + resolved_stages = omni.engine.stage_configs + assert len(resolved_stages) == 1 + stage_config = resolved_stages[0] + + assert isinstance(stage_config, VllmOmniDiffusionStageConfig) + assert stage_config.stage_id == 0 + assert stage_config.model_stage == "dit" + assert stage_config.final_output_type == "video" + + assert omni.engine.stage_vllm_configs == [None] + od_config = omni.get_diffusion_od_config() + assert od_config is not None + assert od_config.stage_id == 0 + assert od_config.model == MODEL + assert od_config.model_class_name == "WanPipeline" + assert od_config.output_type == "pil" + + @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1}) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index 43a29bf3a..366e640f6 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture): assert msg.rpc_id == "abort-rpc-1" assert not task.done() - rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True)) + partial_output = SimpleNamespace(request_id="req-a", finished=True) + rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output])) result = await task - assert result == [] + assert result == [partial_output] try: asyncio.run(_run()) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 15d31cbbc..637835122 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -12,6 +12,7 @@ import types import pytest from omegaconf import OmegaConf +from vllm_omni.config.omni_config import OmniStageRuntimeConfig from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -21,12 +22,23 @@ from vllm_omni.engine.stage_init_utils import ( build_stage0_input_processor, compute_replica_layout, split_devices_for_replicas, + stage_runtime_env, ) from vllm_omni.engine.stage_runtime import StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch): + env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV" + monkeypatch.delenv(env_key, raising=False) + + with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})): + assert os.environ[env_key] == "typed-value" + + assert env_key not in os.environ + + def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch): engine = object.__new__(AsyncOmniEngine) engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True) @@ -223,6 +235,50 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp assert captured["log_stats"] is True +def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch): + import vllm_omni.engine.omni_engine_base as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.stage_configs = [types.SimpleNamespace()] + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.single_stage_mode = False + engine.async_chunk = False + engine.tokenizer = None + engine._single_stage_id_filter = None + engine._omni_master_address = None + engine._omni_master_port = None + engine._omni_dp_size_local = 1 + engine._omni_heartbeat_timeout = 30.0 + engine._omni_lb_policy = "random" + engine.request_queue = types.SimpleNamespace() + engine._log_stats = False + engine._parallel_stage_init = False + + prompt_transform = object() + client = types.SimpleNamespace( + prompt_transform_func=prompt_transform, + prompt_expand_func=None, + default_sampling_params=types.SimpleNamespace(), + final_output=True, + final_output_type="text", + stage_type="llm", + model_stage="text_encoder", + is_comprehension=True, + ) + pool = types.SimpleNamespace( + stage_client=client, + stage_vllm_config=None, + output_processor=None, + ) + runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None) + monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime) + + engine._initialize_stages(stage_init_timeout=7) + + assert engine.prompt_transform_func is prompt_transform + + def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): stage_cfg = types.SimpleNamespace( stage_id=0, @@ -1715,6 +1771,51 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology(): assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} +def test_inject_kv_stage_info_updates_typed_connector_config(): + from vllm_omni.config.omni_config import ( + OmniStageConnectorConfig, + OmniStageDiffusionParallelConfig, + VllmOmniDiffusionStageConfig, + ) + from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + from vllm_omni.entrypoints.utils import inject_omni_kv_config + + stage0 = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + ), + connector_config=OmniStageConnectorConfig( + omni_kv_config={ + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + } + ), + parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4), + ) + stage1 = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=1, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + ), + connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}), + parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2), + ) + + inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1") + inject_kv_stage_info(stage0, 0, [stage0, stage1]) + + assert stage0.connector_config.omni_kv_config["stage_id"] == 0 + assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"} + assert stage0.connector_config.omni_kv_config["engine_input_source"] == [] + assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2} + + def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch): """ROCm default attention logic only applies to LLM stages, not diffusion.""" from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata @@ -1852,3 +1953,80 @@ def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, with pytest.raises(ConfigResolvedError): engine.__init__("test-model") assert engine.async_chunk is expected + + +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed): + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + from vllm_omni.engine.stage_runtime import DistStageRuntime + + if typed: + stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar")) + runtime_cfg = stage.runtime_config + else: + runtime_cfg = types.SimpleNamespace(num_replicas=1) + stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg) + runtime = DistStageRuntime( + stage_configs=[stage], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + single_stage_id_filter=0, + omni_master_address="127.0.0.1", + omni_master_port=12345, + omni_dp_size_local=2, + ) + + runtime._validate_single_stage_mode_replica_constraints() + + assert runtime_cfg.num_replicas == 2 + + +@pytest.mark.parametrize("num_replicas", [1, 2, 3]) +def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas): + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.engine import stage_runtime as runtime_module + + stage = StageConfigFactory.create_typed_default_diffusion( + "generic-diffusion", {"model_class_name": "QwenImagePipeline"} + ).stage_configs[0] + devices = ",".join(str(i) for i in range(num_replicas)) + stage.runtime_config.devices = devices + stage.runtime_config.num_replicas = num_replicas + runtime = StageRuntime( + stage_configs=[stage], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + ) + mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) + mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) + mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext()) + client = mocker.Mock() + launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None)) + + # Replanning must not inherit the first replica's narrowed device slice. + for _ in range(2): + counts, device_map = compute_replica_layout([stage]) + plans = runtime._build_logical_stage_init_plans( + omni_transfer_config=None, + replicas_per_stage=counts, + replica_devices_map=device_map, + ) + replicas = plans[0].replicas + assert stage.runtime_config.devices == devices + assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas + for i, plan in enumerate(replicas): + assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config + assert plan.stage_cfg.runtime_config.devices == str(i) + if num_replicas > 1: + assert plan.stage_cfg is not stage + assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client + assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg + assert launch.call_args.kwargs["metadata"] is plan.metadata + assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1) + + assert launch.call_count == 2 * num_replicas diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py index d88d49825..9f89aabca 100644 --- a/tests/engine/test_parallel_stage_init.py +++ b/tests/engine/test_parallel_stage_init.py @@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util(): assert ledgers[0].kv_budget_bytes == int(cap * 0.3) +def test_check_admission_reads_typed_diffusion_cache_config(): + cap = 80 * 1024**3 + diff = _local_diffusion_replica(0, 0, "0", util=None) + diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35) + plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff]) + + ledgers = check_admission( + [plan], + resolve_physical_devices=lambda _r: [0], + device_total_memory=lambda _d: cap, + external_reserve_bytes=0, + safety_margin_bytes=0, + ) + + assert ledgers[0].kv_budget_bytes == int(cap * 0.35) + + def test_check_admission_unresolved_local_raises(): """A local replica invisible to the ledger must fail admission, not skip it: it would otherwise initialize concurrently with unbounded demand.""" diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 41d63af9e..9ce707dbc 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -16,7 +16,8 @@ from pytest_mock import MockerFixture from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( @@ -57,6 +58,10 @@ def _make_llm_plan( vllm_config: Any | None = None, ) -> LogicalStageInitPlan: stage_cfg = _make_stage_cfg(stage_id) + if launch_mode == "remote": + stage_cfg = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"), + ) metadata = SimpleNamespace( stage_id=stage_id, stage_type="llm", @@ -98,6 +103,16 @@ def _make_diffusion_plan( launch_mode: str, ) -> LogicalStageInitPlan: stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion") + if launch_mode == "remote": + stage_cfg = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=stage_id, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ) metadata = SimpleNamespace( stage_id=stage_id, stage_type="diffusion", @@ -960,6 +975,7 @@ class TestSingleStageReplicaInitialization: assert mock_connect.call_args.kwargs["stage_id"] == 1 assert mock_connect.call_args.kwargs["replica_id"] == 0 assert client_kwargs["log_stats"] is True + assert client_kwargs["metadata"].model_stage == "thinker" assert events == ["enter", "exit", "attach"] def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): @@ -1125,11 +1141,9 @@ class TestSingleStageReplicaInitialization: return _ctx() - remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0] sentinel_client = SimpleNamespace() - mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata) mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect) mock_from_addresses = mocker.patch( "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", @@ -1146,6 +1160,7 @@ class TestSingleStageReplicaInitialization: replica_id=0, ) mock_from_addresses.assert_called_once() + assert mock_from_addresses.call_args.args[0].final_output_type == "image" def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture): import vllm_omni.engine.stage_runtime as runtime_mod diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py index 0473dc341..beb66dfba 100644 --- a/tests/engine/test_stage_device_layout.py +++ b/tests/engine/test_stage_device_layout.py @@ -156,6 +156,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected): assert replica_devices_map == {0: expected} +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize("local_dp", [None, 0, 1, 2]) +@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"]) +def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices): + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig + from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices + + engine_args = { + "tensor_parallel_size": 2, + "data_parallel_size": 4, + "data_parallel_size_local": local_dp, + "pipeline_parallel_size": 2, + } + width = 2 * max(1, local_dp if local_dp is not None else 4) * 2 + devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1))) + if typed: + pipeline = PipelineConfig( + model_type="test", + model_arch="TestModel", + stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),), + ) + stage = VllmOmniConfig.from_pipeline_config( + pipeline, + cli_overrides={ + f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items() + }, + ).stage_configs[0] + else: + stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args) + expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)] + + assert get_stage_devices_per_replica(stage) == width + assert compute_replica_layout([stage]) == ([2], {0: expected}) + assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected + + def test_non_tp_layout_error_uses_generic_guidance(): """A PP mismatch must not be reported as a top-level TP broadcast.""" stage = _stage(0, devices="0") diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 7ac5b6b69..6650c8e45 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -95,6 +95,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( "num_weight_load_threads", "omni_kv_config", "parallel_config", + "pooling_output_decoder", "silence_ban_frames", "subtalker_sampling_params", "task_type", @@ -366,6 +367,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): assert inherited_defaults.isdisjoint(engine_args) +def test_typed_llm_projection_omits_diffusion_only_and_process_only_defaults(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + assert { + "enable_multithread_weight_load", + "num_weight_load_threads", + "disable_autocast", + "log_level", + }.isdisjoint(engine_args) + + def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch): monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 37e2dc430..6e28bb5b1 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -64,6 +64,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # TODO: extend this set and the expected census if @router.head / @router.options appear. _HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"} + +def test_profiler_endpoints_detect_typed_stage_config(): + stage = SimpleNamespace(profiler_config=SimpleNamespace(profiler="torch")) + + assert api_server._should_enable_profiler_endpoints([stage]) + + # Exact Omni router census on main. Update intentionally when routes change. _EXPECTED_ROUTER_ROUTES = { ("POST", "/v1/chat/completions"), diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py index 87524825e..141dd8cac 100644 --- a/tests/entrypoints/openai_api/test_audex_serving_guards.py +++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py @@ -15,6 +15,7 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.entrypoints.openai.serving_speech import ( _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES, OmniOpenAIServingSpeech, @@ -251,9 +252,28 @@ class TestAudexStageDetection: stage = serving._find_tts_stage() assert stage is not None and stage.engine_args.model_stage == "audex_thinker" + def test_typed_pipeline_detected_without_legacy_engine_args(self): + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + stages = [ + SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage), + model_config=SimpleNamespace(model_arch=None), + worker_type="generation", + ) + for index, model_stage in enumerate(("audex_thinker", "audex_code2wav")) + ] + serving.engine_client = SimpleNamespace(stage_configs=stages) + + stage = serving._find_tts_stage() + serving._tts_stage = stage + + assert stage is stages[0] + assert serving._detect_tts_model_type() == "audex" + def test_guided_injection_sets_guided_temperature(): """Guided requests decode at the measured CFG temperature (0.05).""" + serving = TestAudexCfgInjection()._serving_with_fake_tokenizer() cond = TestAudexCfgInjection()._cond_prompt() params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1) @@ -265,6 +285,31 @@ def test_guided_injection_sets_guided_temperature(): assert out.temperature == 0.1 +def test_typed_full_pipeline_uses_full_checkpoint_tokenizer(mocker): + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + serving._tts_stage = SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="audex_omni"), + model_config=SimpleNamespace(model_arch=None), + worker_type="ar", + ) + serving.engine_client = SimpleNamespace( + model_config=SimpleNamespace(model="/models/audex"), + ) + ensure_snapshot = mocker.patch( + "vllm_omni.model_executor.models.audex.checkpoint.ensure_audex_snapshot", + return_value="/snapshots/audex", + ) + from_pretrained = mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=object()) + + adapter = AudexAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client)) + adapter._get_tokenizer() + + ensure_snapshot.assert_called_once_with("/models/audex", profile="full") + from_pretrained.assert_called_once_with("/snapshots/audex/checkpoint_folder_full") + + # ---------------------------------------------------------------- TTA serving surface from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter # noqa: E402 diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 46ff0f822..d274eadda 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -138,6 +138,36 @@ def test_policy_server_config_reads_stage_config_model_config(): assert serving.policy_server_config.to_dict() == policy_config +@pytest.mark.parametrize("inline", [False, True], ids=["out_of_process", "inline"]) +def test_policy_server_config_reads_typed_runtime(inline): + from tests.helpers.stage_config import get_deploy_config_path + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE + from vllm_omni.entrypoints.async_omni import AsyncOmni + + typed_stage = VllmOmniConfig.from_pipeline_config( + PI0_PIPELINE, deploy_config_path=get_deploy_config_path("pi0.yaml") + ).stage_configs[0] + policy_config = typed_stage.diffusion_config.model_config["policy_server_config"] + stage_client = SimpleNamespace(stage_type="diffusion") + if inline: + # Worker-enriched inline metadata remains authoritative over the + # startup stage config; only the process client needs the fallback. + policy_config = {**policy_config, "worker_metadata": "loaded"} + stage_client.od_config = SimpleNamespace(model_config={"policy_server_config": policy_config}) + engine_client = AsyncOmni.__new__(AsyncOmni) + engine_client.engine = SimpleNamespace( + stage_clients=[stage_client], + stage_configs=[typed_stage], + stage_vllm_configs=[None], + get_diffusion_od_config=lambda: SimpleNamespace(model_class_name="Pi0Pipeline"), + ) + + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client) + + assert serving.policy_server_config.to_dict() == policy_config + + def test_policy_server_config_reads_omegaconf_stage_config(): engine_client = SimpleNamespace( get_diffusion_od_config=lambda: None, diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py index 59d070a39..417f457d3 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py @@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @dataclass class _EngineClientStub: - stage_configs: list[dict[str, object]] + stage_configs: list[object] @dataclass @@ -149,6 +149,37 @@ async def test_chat_reference_audio_reaches_code2wav_transport( assert code2wav_payload.meta.ref_audio_sr == 16000 +@pytest.mark.asyncio +async def test_typed_stage_enables_reference_audio( + serving_chat: OmniOpenAIServingChat, + mocker: MockerFixture, +) -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + reference_waveform = np.array([0.25, -0.5], dtype=np.float32) + serving_chat.engine_client = _EngineClientStub( + [ + SimpleNamespace( + model_config=SimpleNamespace(model_arch="MiniCPMO45OmniForConditionalGeneration"), + model_stage="llm", + ) + ] + ) + serving_chat.model_config = _ModelConfigStub() + engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]} + media_connector = mocker.patch( + "vllm_omni.entrypoints.openai.serving_chat.MediaConnector", + ).return_value + media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000)) + + await serving_chat._attach_minicpmo45_reference_audio( + engine_prompt, + _chat_request(reference_audio_url), + ) + + media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url) + assert MINICPMO45_REFERENCE_AUDIO_KEY in cast(dict[str, object], engine_prompt) + + @pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"]) @pytest.mark.asyncio async def test_reference_audio_accepts_request_extra_mappings( diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py index 1a3a4f22e..5ce129d15 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion. Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that @@ -33,6 +34,13 @@ def _stage(model_arch: str, model_stage: str) -> SimpleNamespace: return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage)) +def _typed_stage(model_arch: str, model_stage: str) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace(model_arch=model_arch), + model_stage=model_stage, + ) + + def _minicpmo45_stages(): return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")] @@ -50,6 +58,20 @@ def test_thinker_forced_final_only_talker_stays_delta(serving_chat): assert out[1].output_kind is RequestOutputKind.DELTA +def test_typed_stages_force_thinker_final_only(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + _typed_stage(MINICPMO45_ARCH, "llm"), + _typed_stage(MINICPMO45_ARCH, "tts"), + ] + ) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"]) + + assert out[0].output_kind is RequestOutputKind.FINAL_ONLY + assert out[1].output_kind is RequestOutputKind.DELTA + + @pytest.mark.parametrize("modalities", [None, [], ["text"]]) def test_non_audio_request_is_untouched(serving_chat, modalities): serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages()) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index f8dfaafb6..5ac895d86 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -27,6 +27,7 @@ from pydantic import ValidationError from pytest_mock import MockerFixture from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler from vllm_omni.diffusion.sched.step_scheduler import StepScheduler @@ -1039,6 +1040,38 @@ class TestTTSMethods: asyncio.run(server._prepare_speech_generation(request)) server.shutdown() + def test_prepare_speech_rejects_typed_non_tts_omni_model(self, mocker: MockerFixture): + """Typed Qwen3-Omni stages must hit the same talker guard as legacy stages.""" + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.tts_max_instructions_length = None + mock_engine_client.stage_configs = [ + SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage), + model_config=SimpleNamespace(model_arch=None), + worker_type=worker_type, + ) + for index, (model_stage, worker_type) in enumerate( + ( + ("thinker", "ar"), + ("talker", "ar"), + ("code2wav", "generation"), + ) + ) + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + request = OpenAICreateSpeechRequest(input="Hello world") + with pytest.raises(ValueError, match="only supported for dedicated TTS models"): + asyncio.run(server._prepare_speech_generation(request)) + server.shutdown() + def test_estimate_prompt_len_fallback(self, speech_server): """Test prompt length estimation falls back to 2048 when model is unavailable.""" tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]} diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index e0a165a55..500feab35 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -22,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( detect_tts_model_type, resolve_adapter, ) +from vllm_omni.entrypoints.openai.tts_adapters.base import resolve_stage_model_path from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import ( @@ -99,6 +100,31 @@ def test_resolve_unknown_returns_none(): assert resolve_adapter(None) is None +def test_stage_model_path_prefers_typed_override(): + engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model")), + SimpleNamespace( + model_config=SimpleNamespace(model="typed-stage-model"), + ), + ], + model="served-model", + ) + + assert resolve_stage_model_path(engine_client) == "typed-stage-model" + + +def test_stage_model_path_falls_back_to_legacy_then_served_model(): + legacy_client = SimpleNamespace( + stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model"))], + model="served-model", + ) + served_client = SimpleNamespace(stage_configs=[SimpleNamespace()], model="served-model") + + assert resolve_stage_model_path(legacy_client) == "legacy-stage-model" + assert resolve_stage_model_path(served_client) == "served-model" + + def test_voxcpm2_resolves(): """VoxCPM2 (the served ``latent_generator`` model) resolves cleanly. diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 2de5c4056..0b44cabdf 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -1134,6 +1134,32 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal assert handler.supports_mixed_reference_inputs +def test_typed_stage_drives_video_capability_checks(): + from vllm_omni.config.config_factory import StageConfigFactory + + minimax_stage = StageConfigFactory.create_typed_default_diffusion( + "minimax-h3", + {"model_class_name": "MiniMaxH3Pipeline"}, + ).stage_configs[0] + minimax_handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)), + model_name="minimax-h3", + stage_configs=[minimax_stage], + ) + assert minimax_handler.supports_mixed_reference_inputs + + cosmos_stage = StageConfigFactory.create_typed_default_diffusion( + "cosmos3", + {"model_class_name": "Cosmos3OmniDiffusersPipeline"}, + ).stage_configs[0] + cosmos_handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)), + model_name="cosmos3", + stage_configs=[cosmos_stage], + ) + assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"}) + + @pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name): handler = test_client.app.state.openai_serving_video @@ -1191,6 +1217,35 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames(): assert spec.keep == "first" +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize( + ("num_frames", "extra_params", "expected"), + [ + (189, {"condition_frame_indexes_vision": [0, 2], "condition_video_keep": "last"}, (9, "last")), + (189, {"condition_frame_indexes_vision": [0, 2]}, (9, "first")), + (5, {"condition_frame_indexes_vision": [0, 20]}, (5, "first")), + (None, {"action_mode": "inverse_dynamics", "action_chunk_size": 16}, (17, "first")), + ], +) +def test_cosmos3_reference_video_decode_policy_with_runtime_configs(typed, num_frames, extra_params, expected): + from vllm_omni.config.config_factory import StageConfigFactory + + if typed: + stages = list( + StageConfigFactory.create_typed_default_diffusion( + "cosmos3", + {"model_class_name": "Cosmos3OmniDiffusersPipeline"}, + ).stage_configs + ) + else: + stages = _cosmos3_stage_configs() + request = VideoGenerationRequest(prompt="Continue this motion.", num_frames=num_frames, extra_params=extra_params) + + spec = _reference_video_decode_spec(request, stages) + + assert (spec.max_frames, spec.keep) == expected + + def test_cosmos3_reference_video_limit_preserves_action_frames(): request = VideoGenerationRequest( prompt="Predict the action.", diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 3b378930a..9b99af712 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +from typing import Any import pytest import torch from pydantic import ValidationError from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig +from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig from vllm_omni.engine import stage_init_utils from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -117,7 +119,7 @@ def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocke trust_remote_code=False, ) - assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True def test_default_stage_rejects_unknown_nested_parallel_config_key(): @@ -239,6 +241,33 @@ def test_default_stage_config_includes_default_sampling_params(): } +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}']) +def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults): + from vllm_omni.config.yaml_util import create_config + from vllm_omni.entrypoints.omni_base import OmniBase + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults} + if typed: + stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0] + metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage) + else: + stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0] + metadata = stage_init_utils.extract_legacy_stage_metadata(stage) + base = OmniBase.__new__(OmniBase) + base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage]) + base.default_sampling_params_list = [metadata.default_sampling_params] + base.sampling_constraints_list = base._get_sampling_constraints_list([stage]) + + assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5 + requested = OmniDiffusionSamplingParams(guidance_scale=2.0) + assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0 + assert requested.guidance_scale == 2.0 + assert base.default_sampling_params_list[0].guidance_scale == 7.5 + assert base.sampling_constraints_list == [{}] + + def test_default_stage_config_includes_diffusion_attention_backend(): """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config.""" stage_cfg = StageConfigFactory.create_default_diffusion( @@ -726,3 +755,106 @@ def test_default_stage_config_includes_quantization_config(): stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0] assert stage_cfg["engine_args"]["quantization_config"] == quantization_config + + +@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"]) +def test_generic_diffusion_uses_canonical_video_output_type(model_class_name): + config = StageConfigFactory.create_typed_default_diffusion( + "generic-video", + {"model_class_name": model_class_name}, + ) + + assert config.stage_configs[0].final_output_type == "video" + + +def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker): + """Generic diffusion reaches runtime as the structured stage itself.""" + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + + resolved = resolve_omni_config( + "generic-diffusion", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={ + "num_gpus": 4, + "tensor_parallel_size": 2, + "default_sampling_params": '{"0": {"guidance_scale": 7.5}}', + }, + stage_overrides=None, + strategy_config_path=None, + ) + + assert resolved.pipeline_config is not None + assert resolved.pipeline_config.model_type == "generic_diffusion" + assert len(resolved.stage_configs) == 1 + stage = resolved.stage_configs[0] + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.model_config.model == "generic-diffusion" + assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5} + assert stage.parallel_config.tensor_parallel_size == 2 + assert stage.parallel_config.data_parallel_size == 2 + assert stage.parallel_config.world_size == 4 + assert stage.runtime_config.devices == "0,1,2,3" + + +def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker): + """Standard runtime resolves and starts the typed stage through the real launcher.""" + from vllm_omni.engine import stage_engine_startup as startup_module + from vllm_omni.engine import stage_runtime as runtime_module + from vllm_omni.engine.stage_runtime import StageRuntime + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + resolved = resolve_omni_config( + "generic-diffusion", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={"num_gpus": 1}, + stage_overrides=None, + strategy_config_path=None, + ) + stage = resolved.stage_configs[0] + launched: dict[str, Any] = {} + client = SimpleNamespace(input_address=None, shutdown=mocker.Mock()) + + mocker.patch.object(runtime_module, "prepare_engine_environment") + mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) + mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) + + def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs): + launched.update( + stage_id=stage_id, + model=model, + stage_config=stage_config, + metadata=metadata, + **kwargs, + ) + return client + + mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage) + + runtime = StageRuntime( + stage_configs=list(resolved.stage_configs), + model="generic-diffusion", + config_path="", + stage_init_timeout=10, + async_chunk=False, + ) + runtime.initialize() + + assert launched["stage_config"] is stage + assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig) + assert launched["metadata"].stage_type == "diffusion" + assert launched["metadata"].model_stage == "diffusion" + assert launched["use_inline"] is True + assert launched["model"] == "generic-diffusion" + assert launched["stage_id"] == 0 + assert runtime.stage_pools[0].clients == [client] diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index ea57d07ab..969bbaf91 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -17,6 +17,8 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import StageConfig from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage @@ -202,6 +204,49 @@ def _make_base(): return obj +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +def test_resolve_sampling_params_list_preserves_stage_constraints(typed): + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + + base = _make_base() + base.engine.num_stages = 1 + base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])] + if typed: + stage = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ) + ) + else: + stage = StageConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ).to_omegaconf() + base.engine.stage_configs = [stage] + base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) + assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}] + + resolved_defaults = base.resolve_sampling_params_list(None) + assert resolved_defaults[0].stop_token_ids == [42] + + caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7]) + resolved = base.resolve_sampling_params_list(caller_params) + + assert resolved[0] is not caller_params + assert resolved[0].seed == 1234 + assert resolved[0].max_tokens == 7 + assert resolved[0].detokenize is False + assert resolved[0].stop_token_ids == [7, 42] + assert 7 in resolved[0]._all_stop_token_ids + assert 42 in resolved[0]._all_stop_token_ids + assert caller_params.detokenize is True + assert caller_params.stop_token_ids == [7] + + def test_resolve_sampling_params_list_merges_required_stop_tokens(): base = _make_base() base.engine.num_stages = 1 @@ -219,22 +264,13 @@ def test_resolve_sampling_params_list_merges_required_stop_tokens(): base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}] - resolved_defaults = base.resolve_sampling_params_list(None) - assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645] +def test_sampling_constraints_are_forwarded_by_typed_stage_configs(): + config = VllmOmniConfig.from_pipeline_config(OMNI_PIPELINES["qwen3_tts"]) - caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704]) + constraints = OmniBase._get_sampling_constraints_list(config.stage_configs) - resolved = base.resolve_sampling_params_list(caller_params) - - assert resolved[0] is not caller_params - assert resolved[0].seed == 1234 - assert resolved[0].max_tokens == 7 - assert resolved[0].detokenize is False - assert resolved[0].stop_token_ids == [100, 151704, 151645] - assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids) - assert caller_params.detokenize is True - assert caller_params.stop_token_ids == [100, 151704] + assert constraints == [dict(stage.stage_pipeline_config.sampling_constraints) for stage in config.stage_configs] @pytest.mark.parametrize("use_defaults", [False, True]) diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index 450e1be39..3cfa0f0e5 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -32,15 +32,18 @@ def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch): stage_overrides=None, strategy_config_path=None, ) - engine_args = resolved.stage_configs[0].engine_args + stage = resolved.stage_configs[0] assert resolved.config_path is not None assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml") - assert engine_args.model_class_name == "DreamZeroPipeline" - assert engine_args.model_config.policy_server_config.action_space == "joint_position" + assert stage.diffusion_config.model_class_name == "DreamZeroPipeline" + assert stage.diffusion_config.model_config["policy_server_config"]["action_space"] == "joint_position" def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch): + def _get_hf_config(path, _model): + return None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]} + monkeypatch.setattr( "vllm.transformers_utils.config.get_hf_file_to_dict", lambda path, _model, **_kwargs: ( diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 1f3b58170..8d3766df0 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -7,10 +7,12 @@ from __future__ import annotations import argparse from types import SimpleNamespace +from typing import Any import pytest from pytest_mock import MockerFixture +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.config.resolver import OmniConfigResolution from vllm_omni.entrypoints.cli.serve import ( OmniServeCommand, @@ -334,6 +336,26 @@ def test_parse_stage_overrides_invalid_json_raises() -> None: assert f"Got: {bad!r}" in message +@pytest.mark.parametrize( + "payload", + ['{"abc": {}}', '{"-1": {}}', '{"1.5": {}}', '{"\uff10": {}}'], +) +def test_parse_stage_overrides_rejects_invalid_stage_ids(payload: str) -> None: + with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer stage ids"): + _parse_stage_overrides(payload) + + +def test_parse_stage_overrides_preserves_arbitrary_override_fields() -> None: + parsed = _parse_stage_overrides( + '{"0": {"extras": {"ltx2_use_conv_vae": true}, "kv_cache_dtype": "fp8", "typo_field_xyz": 1}}' + ) + assert parsed["0"] == { + "extras": {"ltx2_use_conv_vae": True}, + "kv_cache_dtype": "fp8", + "typo_field_xyz": 1, + } + + def test_parse_stage_overrides_rejects_non_dict_top_level() -> None: """Top-level must be a JSON object (dict). A list, scalar, or non-dict mapping is rejected with a ValueError naming ``--stage-overrides`` and @@ -706,6 +728,77 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) assert manager_kwargs["omni_replica_id"] == 0 +def test_run_headless_generic_diffusion_launches_structured_stage(mocker: MockerFixture) -> None: + """Headless resolution starts the typed stage through the real group launcher.""" + from vllm_omni.engine import stage_engine_startup as startup_module + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch.object(startup_module.stage_init_utils, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object( + startup_module.initialization, + "resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + captured: dict[str, Any] = {} + od_config = SimpleNamespace() + + def _build_diffusion_config(model, stage_config, metadata): + captured.update(model=model, stage_config=stage_config, metadata=metadata) + return od_config + + def _launch_replica_group(**kwargs): + captured.update(group_kwargs=kwargs) + captured["manager"] = kwargs["launch_one"](0) + + def _launch_replica(**kwargs): + captured.update(replica_kwargs=kwargs) + return SimpleNamespace(exitcode=None) + + mocker.patch.object(startup_module.stage_init_utils, "build_diffusion_config", side_effect=_build_diffusion_config) + mocker.patch.object(startup_module, "launch_headless_replica_group", side_effect=_launch_replica_group) + mocker.patch.object(startup_module, "launch_headless_diffusion_replica", side_effect=_launch_replica) + + explicit_keys = frozenset( + { + "model", + "stage_id", + "omni_master_address", + "omni_master_port", + "worker_backend", + "model_class_name", + "num_gpus", + } + ) + args = _make_headless_args( + explicit_keys=explicit_keys, + model="generic-diffusion", + model_class_name="FakeDiffusionPipeline", + num_gpus=1, + ) + + run_headless(args) + + stage = captured["stage_config"] + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert captured["metadata"].stage_type == "diffusion" + assert captured["metadata"].model_stage == "diffusion" + assert captured["model"] == "generic-diffusion" + assert captured["group_kwargs"]["stage_id"] == 0 + assert captured["group_kwargs"]["omni_dp_size_local"] == 1 + assert captured["group_kwargs"]["per_replica_devices"] == ["0"] + assert captured["replica_kwargs"]["stage_config"] is stage + assert captured["replica_kwargs"]["stage_id"] == 0 + assert captured["replica_kwargs"]["omni_master_address"] == "127.0.0.1" + assert captured["replica_kwargs"]["omni_master_port"] == 26000 + assert captured["replica_kwargs"]["od_config"] is od_config + assert stage.runtime_config.devices == "0" + + def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None: """A diffusion replica that exits with a non-zero code must surface as a RuntimeError from ``run_headless`` (the head needs the signal to roll diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 680bb110c..e7f94ffef 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -13,6 +13,14 @@ import torch from pytest_mock import MockerFixture from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm_omni.config.composable_parallel import ( + Broadcast, + FanInByStage, + MeshAxisSpec, + RouteByStage, + StrategySpec, + TakeRank, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.resolver import ( OmniConfigResolution, @@ -21,7 +29,6 @@ from vllm_omni.config.resolver import ( resolve_omni_config, ) from vllm_omni.config.stage_config import PipelineConfig -from vllm_omni.config.yaml_util import create_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.utils import ( @@ -356,16 +363,12 @@ class TestResolveOmniConfig: ) assert resolved.config_path is None assert len(resolved.stage_configs) == 1 - engine_args = resolved.stage_configs[0]["engine_args"] - assert "dtype" in engine_args - assert engine_args["engine_backend"] == engine_backend - assert engine_args["revision"] == "pinned-revision" + assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32 + assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend + assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision" def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture): - mocker.patch( - "vllm_omni.config.resolver.StageConfigFactory.create_from_model", - return_value=None, - ) + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) resolve_model_class = mocker.patch( "vllm_omni.config.resolver.resolve_model_class_name", return_value="WanImageToVideoPipeline", @@ -374,59 +377,98 @@ class TestResolveOmniConfig: "vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs", return_value={"WanImageToVideoPipeline"}, ) + load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls") resolved = resolve_omni_config( "/models/Wan2.2-I2V", trust_remote_code=False, deploy_config_path=None, - cli_overrides={ - "diffusion_load_format": "diffusers", - "revision": "pinned-revision", - }, + cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"}, stage_overrides=None, strategy_config_path=None, ) + resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision") + load_model_class.assert_not_called() stage = resolved.stage_configs[0] - resolve_model_class.assert_called_once_with( - "/models/Wan2.2-I2V", - "diffusers", - "pinned-revision", - ) - assert stage.engine_args.model_class_name == "WanImageToVideoPipeline" + assert stage.diffusion_config.model_class_name == "WanImageToVideoPipeline" + assert stage.diffusion_config.revision == "pinned-revision" assert stage.final_output_type == "video" + def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "LTX2Pipeline"), + ) + resolved = resolve_omni_config( + "/models/LTX-2.5-Diffusers", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={ + "extras": {"ltx2_use_conv_vae": False, "keep": "global"}, + "tensor_parallel_size": 1, + "diffusion_streaming_output": True, + }, + stage_overrides={ + "0": { + "extras": {"ltx2_use_conv_vae": True}, + "tensor_parallel_size": 2, + "devices": "2,3", + "engine_backend": "custom.backend", + }, + }, + strategy_config_path=None, + ) + stage = resolved.stage_by_id(0) + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers") + + assert stage.runtime_config.devices == "2,3" + assert stage.parallel_config.world_size == 2 + assert engine_args["engine_backend"] == "custom.backend" + assert engine_args["extras"]["ltx2_use_conv_vae"] is True + assert engine_args["extras"]["keep"] == "global" + assert engine_args["streaming_output"] is True + def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture): endpoint_restriction = SimpleNamespace(name="chat") + typed_stage = SimpleNamespace(stage_id=1) structured_config = SimpleNamespace( - orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"), + orchestrator_config=SimpleNamespace( + deploy_config_path="/resolved/deploy.yaml", + omni_lb_policy="round_robin", + ), pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)), + stage_configs=(typed_stage,), + strategy_omni_lb_policy="round-robin", ) - runtime_stage = create_config( - { - "stage_id": 1, - "runtime": {"devices": "1,2,3", "num_replicas": 3}, - "engine_args": {"model": "dummy-model"}, - } - ) - legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage) create_structured = mocker.patch( "vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=structured_config, ) create_legacy = mocker.patch( - "vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry", - return_value=SimpleNamespace( - stage_configs=[legacy_stage], - pipeline_config=structured_config.pipeline_config, - omni_lb_policy="round_robin", - ), + "vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model", ) - strategy_specs = {"stage_1": {"dp": 3}} + strategy_specs = { + "stage_1": [ + StrategySpec( + "stage_replica", + MeshAxisSpec("stage_replica", 2), + RouteByStage("round_robin"), + FanInByStage(), + ) + ] + } load_strategy = mocker.patch( "vllm_omni.config.resolver._load_strategy_specs", return_value=strategy_specs, ) + apply_strategy = mocker.patch( + "vllm_omni.config.composable_parallel.apply_strategy_specs", + return_value=SimpleNamespace(omni_lb_policy="round_robin"), + ) resolved = resolve_omni_config( "dummy-model", @@ -447,31 +489,63 @@ class TestResolveOmniConfig: trust_remote_code=None, cli_overrides=expected_overrides, deploy_config_path="deploy.yaml", - ) - create_legacy.assert_called_once_with( - structured_config.pipeline_config, - expected_overrides, - "/resolved/deploy.yaml", strategy_specs=strategy_specs, ) + create_legacy.assert_not_called() load_strategy.assert_called_once_with("strategy.yaml") + apply_strategy.assert_not_called() assert resolved.config_path == "/resolved/deploy.yaml" assert resolved.pipeline_config is structured_config.pipeline_config - assert resolved.omni_lb_policy == "round_robin" + assert resolved.omni_lb_policy == "round-robin" assert resolved.endpoint_restrictions == (endpoint_restriction,) - assert resolved.stage_configs == (runtime_stage,) + assert resolved.stage_configs == (typed_stage,) - def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture): - pipeline = OMNI_PIPELINES["qwen3_tts"] - assert isinstance(pipeline, PipelineConfig) + def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture): structured_config = SimpleNamespace( - orchestrator_config=SimpleNamespace(deploy_config_path=None), - pipeline_config=pipeline, + orchestrator_config=SimpleNamespace( + deploy_config_path="/resolved/deploy.yaml", + omni_lb_policy="round-robin", + ), + pipeline_config=SimpleNamespace(endpoint_restrictions=()), + stage_configs=(SimpleNamespace(stage_id=0),), + strategy_omni_lb_policy=None, ) mocker.patch( "vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=structured_config, ) + mocker.patch( + "vllm_omni.config.resolver._load_strategy_specs", + return_value={ + "thinker": [ + StrategySpec( + "tp", + MeshAxisSpec("tp", 2), + Broadcast(), + TakeRank(), + ) + ] + }, + ) + + resolved = resolve_omni_config( + "dummy-model", + trust_remote_code=True, + deploy_config_path="deploy.yaml", + cli_overrides={"omni_lb_policy": "round-robin"}, + stage_overrides=None, + strategy_config_path="strategy.yaml", + ) + + assert resolved.omni_lb_policy is None + + def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture): + pipeline = OMNI_PIPELINES["qwen3_tts"] + assert isinstance(pipeline, PipelineConfig) + mocker.patch( + "vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config", + return_value=pipeline, + ) resolved = resolve_omni_config( "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", @@ -483,9 +557,9 @@ class TestResolveOmniConfig: ) assert resolved.pipeline_config is not None - pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages] - runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs] - assert pipeline_stage_ids == runtime_stage_ids + assert [stage.stage_id for stage in resolved.pipeline_config.stages] == [ + stage.stage_id for stage in resolved.stage_configs + ] assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner" assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1 diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py index 879c7aec3..acfb9e229 100644 --- a/vllm_omni/config/composable_parallel/apply.py +++ b/vllm_omni/config/composable_parallel/apply.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Apply declarative parallel strategies onto merged stage configs. This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses @@ -198,10 +198,41 @@ def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: Ro def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None: + declared = set(cfg.l1_owners.keys()) + if hasattr(stage, "parallel_config") and hasattr(stage, "runtime_config"): + parallel = stage.parallel_config + for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): + if kind in declared: + existing = getattr(parallel, field_name, None) + value = getattr(cfg, field_name) + if existing not in (None, 1, value): + _fail( + f"role {role!r}: strategy derives {field_name}={value} " + f"but structured config already set {existing}" + ) + setattr(parallel, field_name, value) + if "ep" in declared and cfg.enable_expert_parallel: + setattr(parallel, "enable_expert_parallel", True) + if "stage_replica" in declared: + existing = stage.runtime_config.num_replicas + if existing not in (None, 1, cfg.stage_replica_size): + _fail( + f"role {role!r}: strategy derives " + f"num_replicas={cfg.stage_replica_size} but structured config already set {existing}" + ) + stage.runtime_config.num_replicas = cfg.stage_replica_size + check_device_layout( + stage.runtime_config.devices, + tensor_parallel_size=cfg.tensor_parallel_size, + data_parallel_size=cfg.data_parallel_size, + pipeline_parallel_size=cfg.pipeline_parallel_size, + num_replicas=int(stage.runtime_config.num_replicas or 1), + role=role, + ) + return + engine_args = stage.yaml_engine_args runtime = stage.yaml_runtime - - declared = set(cfg.l1_owners.keys()) for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): if kind in declared: _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role) @@ -209,7 +240,6 @@ def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> No _set_explicit(engine_args, "enable_expert_parallel", True, role=role) if "stage_replica" in declared: _set_num_replicas(runtime, cfg.stage_replica_size, role=role) - _check_devices(runtime, cfg, role=role) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 90108e7df..94decce9a 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -25,6 +25,9 @@ from vllm_omni.config.stage_config import ( DeployConfig, PipelineConfig, StageConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, StageType, build_stage_runtime_overrides, load_deploy_config, @@ -373,6 +376,7 @@ class StageConfigFactory: trust_remote_code: bool | None, cli_overrides: dict[str, Any], deploy_config_path: str | None, + strategy_specs: Mapping[Any, Any] | None = None, ) -> VllmOmniConfig | None: """Build the structured Omni config for a model/deploy pair.""" user_deploy_config = cls._load_user_deploy_config(deploy_config_path) @@ -387,6 +391,20 @@ class StageConfigFactory: if pipeline_cfg is None: return None + # Keep the typed representation's topology in lockstep with the + # compatibility launch representation. The aligner is a real extra + # stage, not an engine-only setting, so it must be injected before + # either representation is built. + if user_deploy_config is None: + if pipeline_cfg.default_deploy_config_name is not None: + default_deploy_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name + user_deploy_config = load_deploy_config(default_deploy_path) + deploy_config_path = str(default_deploy_path) + else: + user_deploy_config = DeployConfig() + from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage + + pipeline_cfg, user_deploy_config = inject_forced_aligner_stage(pipeline_cfg, user_deploy_config, cli_overrides) registry_cli_overrides = with_trust_remote_code_override( {**cli_overrides, "model": model}, trust_remote_code, @@ -396,6 +414,7 @@ class StageConfigFactory: user_deploy_config=user_deploy_config, deploy_config_path=deploy_config_path, cli_overrides=registry_cli_overrides, + strategy_specs=strategy_specs, ) @classmethod @@ -621,13 +640,11 @@ class StageConfigFactory: ) @classmethod - def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Build the temporary runtime ABI for a generic diffusion stage. - - The terminal diffusion config owns engine defaults and normalization; - this compatibility builder only adds Omni stage topology, request - defaults, and device placement. - """ + def _normalize_default_diffusion( + cls, + kwargs: dict[str, Any], + ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]: + """Normalize inputs shared by typed and compatibility diffusion builders.""" raw_sampling_params = kwargs.get("default_sampling_params") if isinstance(raw_sampling_params, str): try: @@ -643,8 +660,6 @@ class StageConfigFactory: if kwargs.get("num_gpus") is not None: parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"])) engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs) - engine_args["parallel_config"] = asdict(parallel_config) - engine_args["model_stage"] = "diffusion" extras = dict(engine_args.get("extras") or {}) for key, default in ( @@ -665,6 +680,21 @@ class StageConfigFactory: final_output_type, model_class_name, ) + return engine_args, parallel_config, default_sampling_params, final_output_type + + @classmethod + def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Build the temporary runtime ABI for a generic diffusion stage. + + The terminal diffusion config owns engine defaults and normalization; + this compatibility builder only adds Omni stage topology, request + defaults, and device placement. + """ + engine_args, parallel_config, default_sampling_params, final_output_type = cls._normalize_default_diffusion( + kwargs + ) + engine_args["parallel_config"] = asdict(parallel_config) + engine_args["model_stage"] = "diffusion" return [ { @@ -681,6 +711,50 @@ class StageConfigFactory: } ] + @classmethod + def create_typed_default_diffusion( + cls, + model: str, + kwargs: dict[str, Any], + ) -> VllmOmniConfig: + """Build generic diffusion directly into the structured runtime config.""" + engine_overrides, parallel_config, default_sampling_params, final_output_type = ( + cls._normalize_default_diffusion(kwargs) + ) + engine_overrides.update(asdict(parallel_config)) + model_class_name = engine_overrides.get("model_class_name") + pipeline = PipelineConfig( + model_type="generic_diffusion", + model_arch=str(model_class_name or ""), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type=final_output_type, + ), + ), + ) + + # These values have already been normalized for this diffusion stage. + # Scope them explicitly so diffusion-only fields such as engine_backend + # and extras are not filtered by the global LLM CLI argument surface. + stage_overrides = { + f"stage_0_{key}": value for key, value in engine_overrides.items() if key not in {"model", "stage_id"} + } + stage_overrides["stage_0_devices"] = ( + kwargs.get("stage_0_devices") + or kwargs.get("devices") + or ",".join(str(i) for i in range(parallel_config.world_size)) + ) + # Caller defaults belong to deployment, not immutable pipeline constraints: + # explicit request sampling parameters must still be able to override them. + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, default_sampling_params=default_sampling_params)]) + return VllmOmniConfig.from_pipeline_config( + pipeline, user_deploy_config=deploy, cli_overrides={"model": model, **stage_overrides} + ) + @classmethod def _merge_cli_overrides( cls, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c13935489..1407615ab 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -29,7 +29,9 @@ from vllm.config import ProfilerConfig as VllmProfilerConfig from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.config.utils import config from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.pooling_params import PoolingParams from vllm_omni.config.stage_config import ( _DEPLOY_DIR, @@ -42,11 +44,13 @@ from vllm_omni.config.stage_config import ( StagePipelineConfig, StageType, _apply_platform_overrides, + _get_recursively_merged_dict, _resolve_scheduler, _scheduler_path, _select_processor_funcs, build_stage_runtime_overrides, load_deploy_config, + merge_pipeline_deploy, merge_sampling_constraints, normalize_pipeline_cli_overrides, reconcile_diffusion_attention_overrides, @@ -55,6 +59,8 @@ from vllm_omni.config.stage_config import ( ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +logger = init_logger(__name__) + _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = { StageExecutionType.LLM_AR: (StageType.LLM, "ar"), StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), @@ -186,6 +192,11 @@ class _ModelEngineOverrides(TypedDict, total=False): enable_prompt_embeds: bool +class _PoolingEngineOverrides(TypedDict, total=False): + runner: str + pooling_output_decoder: str + + class _LoadEngineOverrides(TypedDict, total=False): tokenizer: str download_dir: str @@ -261,6 +272,7 @@ class _StageEngineValues: quantization: _QuantizationEngineOverrides model: _ModelEngineOverrides load: _LoadEngineOverrides + pooling: _PoolingEngineOverrides cache: _CacheEngineOverrides scheduler: _SchedulerEngineOverrides connector: _ConnectorEngineOverrides @@ -476,6 +488,15 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): enable_prompt_embeds: bool | None = None +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class OmniStagePoolingConfig: + """Typed inputs owned by vLLM pooling stages.""" + + runner: str | None = None + pooling_output_decoder: str | None = None + default_pooling_params: PoolingParams | None = None + + @_enforce_keyword_only_init @config(kw_only=True) class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig): @@ -833,12 +854,12 @@ class _DiffusionConfigProjection: from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config - normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) + config_kwargs = normalize_omni_diffusion_kwargs(kwargs) # Validate before stage construction while retaining the raw mapping # needed by dataclass/config serialization across process boundaries. - parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config")) + parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config")) valid_fields = {f.name for f in fields(cast(Any, cls))} - return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) + return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields}) def __post_init__(self) -> None: # Keep diffusion imports lazy so importing vllm_omni.config does not @@ -1145,6 +1166,7 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) +_POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) _DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"}) @@ -1175,6 +1197,7 @@ _LLM_STAGE_ENGINE_FIELDS = ( | _LLM_CACHE_ENGINE_FIELDS | _LLM_SCHEDULER_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS + | _POOLING_ENGINE_FIELDS | {"parallel_config"} ) _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = ( @@ -1308,7 +1331,12 @@ def _stage_engine_values( if topology.execution_type == StageExecutionType.DIFFUSION: # Mirror StageConfig.to_omegaconf so both projections resolve alike. reconcile_diffusion_attention_overrides(engine, stage_cli_overrides) - engine.update(_copy_value(stage_cli_overrides)) + for key, value in stage_cli_overrides.items(): + existing = engine.get(key) + if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, Mapping): + engine[key] = _get_recursively_merged_dict(existing, dict(value)) + else: + engine[key] = _copy_value(value) _validate_stage_engine_override_ownership( topology.stage_id, topology.execution_type, @@ -1334,6 +1362,10 @@ def _stage_engine_values( ), model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)), load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)), + pooling=cast( + _PoolingEngineOverrides, + _select_engine_overrides(engine, _POOLING_ENGINE_FIELDS), + ), cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)), scheduler=cast( _SchedulerEngineOverrides, @@ -1406,6 +1438,7 @@ class BaseVllmOmniStageConfig: cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig) scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig) connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) + pooling_config: OmniStagePoolingConfig = field(default_factory=OmniStagePoolingConfig) runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) compilation_config: VllmCompilationConfig | None = None @@ -1490,6 +1523,10 @@ class BaseVllmOmniStageConfig: def prompt_transform_func(self) -> str | None: return self.stage_pipeline_config.prompt_transform_func + @property + def sampling_constraints(self) -> dict[str, Any]: + return dict(self.stage_pipeline_config.sampling_constraints) + @property def cfg_kv_collect_func(self) -> str | None: return self.stage_pipeline_config.cfg_kv_collect_func @@ -1543,6 +1580,7 @@ def _build_common_stage_config_kwargs( model=model, ), "load_config": _build_load_config(topology, engine.load), + "pooling_config": _build_pooling_config(stage_deploy, engine.pooling), "cache_config": _build_cache_config( deploy, engine.cache, @@ -1756,6 +1794,20 @@ def _build_model_config( ) +def _build_pooling_config( + stage_deploy: StageDeployConfig | None, + engine: _PoolingEngineOverrides, +) -> OmniStagePoolingConfig: + default_pooling_params = None + if stage_deploy is not None and stage_deploy.default_pooling_params: + default_pooling_params = PoolingParams(**dict(stage_deploy.default_pooling_params)) + return cast(Any, OmniStagePoolingConfig)( + runner=_copy_value(engine.get("runner")), + pooling_output_decoder=_copy_value(engine.get("pooling_output_decoder")), + default_pooling_params=default_pooling_params, + ) + + def _build_load_config( topology: StagePipelineConfig, engine: _LoadEngineOverrides, @@ -1934,6 +1986,10 @@ class VllmOmniConfig: pipeline_config: PipelineConfig stage_configs: tuple[StageConfigType, ...] orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig) + # Keep strategy provenance separate from the effective orchestrator value: + # an explicit CLI policy may be present even when the strategy declares no + # stage-replica axis. + strategy_omni_lb_policy: str | None = None def stage_by_id(self, stage_id: int) -> StageConfigType: for stage in self.stage_configs: @@ -1949,6 +2005,7 @@ class VllmOmniConfig: user_deploy_config: DeployConfig | None = None, deploy_config_path: str | None = None, cli_overrides: dict[str, Any] | None = None, + strategy_specs: Mapping[Any, Any] | None = None, ) -> VllmOmniConfig: """Create a structured config from a resolved pipeline and deploy YAML.""" if cli_overrides is None: @@ -1973,6 +2030,62 @@ class VllmOmniConfig: deploy.async_chunk = False _validate_async_chunk_support(pipeline_cfg, deploy) validate_stage_async_chunk_edges(pipeline_cfg, deploy) + + strategy_result = None + if strategy_specs: + from vllm_omni.config.composable_parallel import apply_strategy_specs + + strategy_stages = merge_pipeline_deploy(pipeline_cfg, copy.deepcopy(deploy), {}) + strategy_result = apply_strategy_specs(strategy_stages, strategy_specs) + strategy_overrides: dict[str, Any] = {} + axis_fields = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", + } + for stage_id, parallel in strategy_result.per_stage_config.items(): + declared = set(parallel.l1_owners) + explicit = build_stage_runtime_overrides(stage_id, cli_overrides) + for axis, field_name in axis_fields.items(): + if axis not in declared: + continue + derived = getattr(parallel, field_name) + value = explicit.get(field_name) + if value is None: + value = derived + elif value != derived: + logger.warning( + "[composable_parallel] stage %s: CLI %s=%s overrides the " + "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.", + stage_id, + field_name, + value, + field_name, + derived, + ) + strategy_overrides[f"stage_{stage_id}_{field_name}"] = value + if "ep" in declared and parallel.enable_expert_parallel: + strategy_overrides[f"stage_{stage_id}_enable_expert_parallel"] = bool( + explicit.get("enable_expert_parallel", True) + ) + if "stage_replica" in declared: + derived = parallel.stage_replica_size + value = explicit.get("num_replicas") + if value is None: + value = derived + elif value != derived: + logger.warning( + "[composable_parallel] stage %s: CLI num_replicas=%s overrides the " + "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.", + stage_id, + value, + derived, + ) + strategy_overrides[f"stage_{stage_id}_num_replicas"] = value + cli_overrides = {**strategy_overrides, **cli_overrides} + if strategy_result.omni_lb_policy is not None and cli_overrides.get("omni_lb_policy") is None: + cli_overrides["omni_lb_policy"] = strategy_result.omni_lb_policy + deploy_by_id = {stage.stage_id: stage for stage in deploy.stages} model = cli_overrides.get("model") @@ -1995,6 +2108,20 @@ class VllmOmniConfig: ) for topology in pipeline_cfg.stages ) + if strategy_result is not None: + from vllm_omni.config.composable_parallel import check_device_layout + + by_id = {stage.stage_id: stage for stage in stage_configs} + for stage_id in strategy_result.per_stage_config: + stage = by_id[stage_id] + check_device_layout( + stage.runtime_config.devices, + tensor_parallel_size=stage.parallel_config.tensor_parallel_size, + data_parallel_size=stage.parallel_config.data_parallel_size, + pipeline_parallel_size=stage.parallel_config.pipeline_parallel_size, + num_replicas=stage.runtime_config.num_replicas, + role=stage.model_stage, + ) orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)( deploy_config_path=loaded_deploy_config_path, @@ -2004,6 +2131,7 @@ class VllmOmniConfig: pipeline_config=pipeline_cfg, stage_configs=stage_configs, orchestrator_config=orchestrator_config, + strategy_omni_lb_policy=(strategy_result.omni_lb_policy if strategy_result is not None else None), ) diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py index 4f5ee8a6c..0db436768 100644 --- a/vllm_omni/config/resolver.py +++ b/vllm_omni/config/resolver.py @@ -14,7 +14,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot from vllm_omni.config.endpoint_policy import EndpointRestriction from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.stage_config import PipelineConfig -from vllm_omni.config.yaml_util import create_config from vllm_omni.diffusion.data import resolve_model_class_name from vllm_omni.diffusion.registry import DiffusionModelRegistry from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model @@ -24,19 +23,10 @@ logger = init_logger(__name__) @dataclass(frozen=True) class OmniConfigResolution: - """Migration envelope returned by the production config resolver. - - ``stage_configs`` intentionally carries the current OmegaConf-compatible - runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly. - ``pipeline_config`` is the effective topology after runtime stage injection, - so both fields always describe the same set of stages. - It is not a stable authoring or extension API; new production callers - should resolve through :func:`resolve_omni_config` and must not construct or - merge this compatibility shape themselves. - """ + """Canonical typed result returned by the production config resolver.""" config_path: str | None - stage_configs: tuple[Any, ...] # Temporary StageConfig/OmegaConf bridge. + stage_configs: tuple[Any, ...] pipeline_config: PipelineConfig | None = None omni_lb_policy: str | None = None @@ -176,27 +166,13 @@ def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any] def _build_registered_resolution( structured_config: VllmOmniConfig, - *, - cli_overrides: dict[str, Any], - strategy_config_path: str | None, ) -> OmniConfigResolution: - """Build the temporary runtime view for an already-resolved pipeline.""" - # Runtime consumers have not yet moved to typed per-stage configs. Use the - # factory-owned compatibility bridge instead of reimplementing legacy YAML - # discovery and merging in this resolver. - effective_deploy_path = structured_config.orchestrator_config.deploy_config_path - legacy_resolution = StageConfigFactory._resolve_legacy_from_registry( - structured_config.pipeline_config, - cli_overrides, - effective_deploy_path, - strategy_specs=_load_strategy_specs(strategy_config_path), - ) - + """Return the canonical structured runtime view for a resolved pipeline.""" return OmniConfigResolution( - config_path=effective_deploy_path, - stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs), - pipeline_config=legacy_resolution.pipeline_config, - omni_lb_policy=legacy_resolution.omni_lb_policy, + config_path=structured_config.orchestrator_config.deploy_config_path, + stage_configs=tuple(structured_config.stage_configs), + pipeline_config=structured_config.pipeline_config, + omni_lb_policy=getattr(structured_config, "strategy_omni_lb_policy", None), ) @@ -231,18 +207,16 @@ def resolve_omni_config( registry_overrides = dict(normalized_overrides) _flatten_stage_overrides(registry_overrides, stage_overrides) + strategy_specs = _load_strategy_specs(strategy_config_path) structured_config = StageConfigFactory.create_from_model( model, trust_remote_code=trust_remote_code, cli_overrides=registry_overrides, deploy_config_path=deploy_config_path, + strategy_specs=strategy_specs, ) if structured_config is not None: - return _build_registered_resolution( - structured_config, - cli_overrides=registry_overrides, - strategy_config_path=strategy_config_path, - ) + return _build_registered_resolution(structured_config) _apply_generic_stage_overrides(normalized_overrides, stage_overrides) supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides) @@ -252,11 +226,12 @@ def resolve_omni_config( ) if model_class_name is not None: normalized_overrides.setdefault("model_class_name", model_class_name) - default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides) + default_config = StageConfigFactory.create_typed_default_diffusion(model, normalized_overrides) return OmniConfigResolution( config_path=deploy_config_path, - stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))), + stage_configs=tuple(default_config.stage_configs), + pipeline_config=default_config.pipeline_config, ) diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 85310e9dd..d615ebca0 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -272,7 +272,14 @@ class OmniEngineBase: self.num_stages = len(self.stage_configs) self.async_chunk = any( - bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs + bool( + getattr( + getattr(stage, "connector_config", None), + "async_chunk", + getattr(getattr(stage, "engine_args", None), "async_chunk", False), + ) + ) + for stage in self.stage_configs ) self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers diff --git a/vllm_omni/engine/stage_admission.py b/vllm_omni/engine/stage_admission.py index 469787bca..bd57f805d 100644 --- a/vllm_omni/engine/stage_admission.py +++ b/vllm_omni/engine/stage_admission.py @@ -287,8 +287,13 @@ def check_admission( def _diffusion_utilization(replica: Any) -> float | None: - """Best-effort gpu_memory_utilization for a diffusion replica from raw args.""" + """Resolve a diffusion replica's memory budget from typed or legacy config.""" stage_cfg = getattr(replica, "stage_cfg", None) + cache_config = getattr(stage_cfg, "cache_config", None) + util = getattr(cache_config, "gpu_memory_utilization", None) + if util is not None: + return float(util) + engine_args = getattr(stage_cfg, "engine_args", None) if isinstance(engine_args, dict): util = engine_args.get("gpu_memory_utilization") diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 71ba17c7c..612c7481f 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -17,6 +17,7 @@ from types import SimpleNamespace from typing import Any import msgspec +import torch import zmq from omegaconf import OmegaConf from vllm.config import VllmConfig @@ -33,6 +34,7 @@ from vllm.v1.engine.utils import ( ) from vllm.v1.executor import Executor +from vllm_omni.config.omni_config import BaseVllmOmniStageConfig from vllm_omni.distributed.omni_connectors.utils import initialization from vllm_omni.engine import stage_init_utils from vllm_omni.engine.stage_init_utils import ( @@ -67,11 +69,13 @@ def _serialize_stage_config(stage_config: Any) -> Any: """Convert a stage config to msgpack-friendly builtins.""" if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)): return stage_config + if isinstance(stage_config, torch.dtype): + return str(stage_config).removeprefix("torch.") if OmegaConf.is_config(stage_config): return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True)) - if dataclasses.is_dataclass(stage_config): + if dataclasses.is_dataclass(stage_config) and not isinstance(stage_config, type): return _serialize_stage_config(dataclasses.asdict(stage_config)) if isinstance(stage_config, dict): @@ -1403,7 +1407,7 @@ def get_headless_replica_devices( omni_dp_size_local: int, ) -> list[str | None]: """Return per-replica device slices for a headless stage.""" - runtime_cfg = getattr(stage_cfg, "runtime", None) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None)) devices_str: str | None = None if runtime_cfg is not None: devices_str = ( @@ -1509,10 +1513,11 @@ def launch_headless_diffusion_replicas( stage_id, ) - # Headless diffusion startup and its downstream helpers still consume the - # legacy StageConfig shape; switch this with the coordinated RFC #4021 - # stage-init cutover. - metadata = stage_init_utils.extract_legacy_stage_metadata(stage_cfg) + metadata = ( + stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage_cfg) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else stage_init_utils.extract_legacy_stage_metadata(stage_cfg) + ) if omni_conn_cfg: inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) # Headless single-stage launch must still infer cross-stage TP topology diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9668dfc99..06f916394 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -402,6 +402,11 @@ def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | Non for stage_cfg in stage_configs: if str(getattr(stage_cfg, "stage_id", None)) not in id_strs: continue + if isinstance(stage_cfg, BaseVllmOmniStageConfig): + try: + return max(1, int(stage_cfg.parallel_config.tensor_parallel_size)) + except (TypeError, ValueError): + return 1 engine_args = getattr(stage_cfg, "engine_args", None) if engine_args is None: return 1 @@ -489,11 +494,15 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ rank mappings automatically. """ try: - engine_args = stage_cfg.engine_args - if hasattr(engine_args, "get"): - omni_kv = engine_args.get("omni_kv_config", None) + connector_config = getattr(stage_cfg, "connector_config", None) + if connector_config is not None: + omni_kv = connector_config.omni_kv_config else: - omni_kv = getattr(engine_args, "omni_kv_config", None) + engine_args = stage_cfg.engine_args + if hasattr(engine_args, "get"): + omni_kv = engine_args.get("omni_kv_config", None) + else: + omni_kv = getattr(engine_args, "omni_kv_config", None) if omni_kv is None: return @@ -504,7 +513,11 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ if "stage_id" not in omni_kv: omni_kv["stage_id"] = stage_id - engine_input_source = getattr(stage_cfg, "engine_input_source", None) + engine_input_source = ( + getattr(stage_cfg, "input_sources", None) + if connector_config is not None + else getattr(stage_cfg, "engine_input_source", None) + ) if engine_input_source is not None: if hasattr(omni_kv, "setdefault"): omni_kv.setdefault("engine_input_source", list(engine_input_source)) @@ -703,17 +716,16 @@ def _resolve_omni_metadata_hook(path: str | None) -> Callable | None: def extract_stage_metadata_from_omni_stage_config( stage_config: BaseVllmOmniStageConfig, ) -> StageMetadata: - """Project one typed stage config into metadata for a future cutover. - - This projection is not used by production startup yet. Current replica - layout, engine-argument, remote-diffusion, and platform setup paths still - require the legacy StageConfig/OmegaConf shape. - """ + """Project one typed stage config into production runtime metadata.""" stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm" - sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams - sampling_params: OmniSamplingParams = sampling_params_cls( - **(stage_config.model_config.default_sampling_params or {}) - ) + pooling_config = stage_config.pooling_config + if stage_type == "llm" and (pooling_config.runner or "").lower() == "pooling": + sampling_params: OmniSamplingParams | PoolingParams = ( + copy.deepcopy(pooling_config.default_pooling_params) or PoolingParams() + ) + else: + sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams + sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) if stage_type == "diffusion": @@ -878,7 +890,9 @@ def _get_local_llm_parallel_sizes( unset. """ if engine_args is None: - engine_args = getattr(stage_cfg, "engine_args", {}) + engine_args = getattr(stage_cfg, "parallel_config", None) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1) pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1) local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None) @@ -889,9 +903,12 @@ def _get_local_llm_parallel_sizes( def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - if engine_args is None: - engine_args = getattr(stage_cfg, "engine_args", {}) if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + typed_parallel_config = getattr(stage_cfg, "parallel_config", None) + if typed_parallel_config is not None: + return max(1, int(typed_parallel_config.world_size)) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) parallel_config = _get_attr_or_item(engine_args, "parallel_config") if parallel_config is None: return 1 @@ -933,7 +950,7 @@ def compute_replica_layout( """ replicas_per_stage: list[int] = [] for stage_cfg in stage_configs: - runtime_cfg = getattr(stage_cfg, "runtime", {}) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {})) num_replicas = int( runtime_cfg.get("num_replicas", 1) if hasattr(runtime_cfg, "get") @@ -948,7 +965,7 @@ def compute_replica_layout( num_replicas = replicas_per_stage[stage_id] if num_replicas <= 1: continue - runtime_cfg = getattr(stage_cfg, "runtime", {}) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {})) devices_str = ( runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) ) @@ -1086,15 +1103,31 @@ def _project_omni_stage_engine_args( diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config) engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config)) + model_excluded_fields = { + "default_sampling_params", + "has_sampling_extra_args", + } + runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"} + if not is_diffusion: + # These values configure OmniDiffusionConfig or its worker process; + # OmniEngineArgs has no matching fields for LLM stages. + model_excluded_fields.update( + { + "disable_autocast", + "enable_multithread_weight_load", + "num_weight_load_threads", + } + ) + runtime_excluded_fields.add("log_level") + for config, excluded_fields in ( ( stage_config.model_config, - frozenset({"default_sampling_params", "has_sampling_extra_args"}), + frozenset(model_excluded_fields), ), ( stage_config.runtime_config, - frozenset({"devices", "num_replicas", "env", "num_gpus"}) - | (frozenset({"additional_config"}) if is_diffusion else frozenset()), + frozenset(runtime_excluded_fields) | (frozenset({"additional_config"}) if is_diffusion else frozenset()), ), ): engine_args.update( @@ -1103,6 +1136,11 @@ def _project_omni_stage_engine_args( exclude=excluded_fields, ) ) + pooling_config = stage_config.pooling_config + if pooling_config.runner is not None: + engine_args["runner"] = pooling_config.runner + if pooling_config.pooling_output_decoder is not None: + engine_args["pooling_output_decoder"] = pooling_config.pooling_output_decoder if is_diffusion: for config, field_map in ( @@ -1330,11 +1368,9 @@ def build_engine_args_dict( stage_connector_spec: dict[str, Any] | None = None, cli_tokenizer: str | None = None, ) -> dict[str, Any]: - """Build engine arguments through the stable production entry point. + """Preserve the legacy engine-argument API for compatibility callers. - Production inputs still use the legacy stage representation. Keep that - compatibility choice behind this function so callers do not bind directly - to a representation-specific implementation. + Typed runtime stages use ``build_engine_args_dict_from_omni_stage_config``. """ return build_legacy_engine_args_dict( stage_config, @@ -1350,13 +1386,7 @@ def build_engine_args_dict_from_omni_stage_config( stage_connector_spec: dict[str, Any] | None = None, cli_tokenizer: str | None = None, ) -> dict[str, Any]: - """Project one typed stage config into backend engine arguments. - - This projection is prepared for the RFC #4021 stage-init cutover. Current - production startup reaches the legacy implementation through - ``build_engine_args_dict`` while strategy and startup-plan inputs still - use the legacy representation. - """ + """Project one typed production stage config into backend engine arguments.""" engine_args_dict = _project_omni_stage_engine_args(stage_config) _apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type) return _finalize_engine_args_dict( @@ -1391,7 +1421,7 @@ def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, An """ from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout - runtime = getattr(stage_config, "runtime", None) + runtime = getattr(stage_config, "runtime_config", getattr(stage_config, "runtime", None)) devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None if devices is None: # No explicit placement -> vLLM assigns devices itself; nothing to check. @@ -1452,10 +1482,12 @@ def build_vllm_config( (vllm_config, executor_class) """ if engine_args_dict is None: - engine_args_dict = build_engine_args_dict( - stage_config, - model, - stage_connector_spec=stage_connector_spec, + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_config, model, stage_connector_spec=stage_connector_spec + ) + if isinstance(stage_config, BaseVllmOmniStageConfig) + else build_engine_args_dict(stage_config, model, stage_connector_spec=stage_connector_spec) ) filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict) @@ -1927,7 +1959,11 @@ def build_diffusion_config( ) -> Any: """Build diffusion config for a stage.""" - engine_args_dict = build_engine_args_dict(stage_cfg, model) + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config(stage_cfg, model) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else build_engine_args_dict(stage_cfg, model) + ) od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict) num_devices_per_stage = od_config.parallel_config.world_size diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 53400df5e..d17e3b08f 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -16,9 +16,9 @@ from dataclasses import dataclass from typing import Any, cast import janus -from omegaconf import OmegaConf from vllm.logger import init_logger +from vllm_omni.config.omni_config import BaseVllmOmniStageConfig from vllm_omni.distributed.omni_connectors.utils.initialization import ( resolve_omni_kv_config_for_stage, ) @@ -49,11 +49,13 @@ from vllm_omni.engine.stage_init_utils import ( _inject_inferred_kv_tp_topology, acquire_device_locks, build_engine_args_dict, + build_engine_args_dict_from_omni_stage_config, build_llm_stage_output_processor, build_vllm_config, compute_replica_layout, device_overlap_group_keys, extract_legacy_stage_metadata, + extract_stage_metadata_from_omni_stage_config, get_stage_connector_spec, inject_kv_stage_info, inject_omni_kv_connector_config, @@ -345,12 +347,12 @@ class StageRuntime: """Build startup plans for every logical stage and replica.""" stage_plans: list[LogicalStageInitPlan] = [] - # RFC #4021 transition boundary: stage planning still relies on legacy - # StageConfig.runtime and StageConfig.engine_args for replica and engine - # setup. Keep metadata extraction on the legacy path until the - # coordinated stage-init cutover. for stage_idx, stage_cfg in enumerate(self._stage_configs): - base_metadata = extract_legacy_stage_metadata(stage_cfg) + base_metadata = ( + extract_stage_metadata_from_omni_stage_config(stage_cfg) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else extract_legacy_stage_metadata(stage_cfg) + ) stage_id = int(base_metadata.stage_id) if stage_id != stage_idx: raise ValueError( @@ -374,14 +376,20 @@ class StageRuntime: executor_class = None engine_args_dict = None if base_metadata.stage_type != "diffusion": - # The stable adapter entry point still receives the same - # legacy stage object as replica planning. Its implementation - # switches only at the coordinated RFC #4021 cutover. - engine_args_dict = build_engine_args_dict( - stage_cfg, - self._model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=self._tokenizer, + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else build_engine_args_dict( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) ) inject_omni_kv_connector_config( engine_args_dict, @@ -401,11 +409,20 @@ class StageRuntime: ) for replica_id in range(num_replicas): - replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg + # Keep the logical stage's device pool intact; each replica owns + # the same config throughout planning, metadata and launch. + replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 else stage_cfg if stage_idx in replica_devices_map: - replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] + devices = replica_devices_map[stage_idx][replica_id] + runtime_cfg = getattr(replica_cfg, "runtime_config", getattr(replica_cfg, "runtime", None)) + if runtime_cfg is not None: + runtime_cfg.devices = devices - replica_metadata = extract_legacy_stage_metadata(replica_cfg) + replica_metadata = ( + extract_stage_metadata_from_omni_stage_config(replica_cfg) + if isinstance(replica_cfg, BaseVllmOmniStageConfig) + else extract_legacy_stage_metadata(replica_cfg) + ) replica_metadata.replica_id = replica_id if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": replica_metadata.runtime_cfg = None @@ -859,17 +876,25 @@ class StageRuntime: raise RuntimeError("Omni KV connector requires source and destination stages") inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs) - engine_args = getattr(plan.stage_cfg, "engine_args", {}) - inline_diffusion = ( - engine_args.get("inline_diffusion", False) - if hasattr(engine_args, "get") - else getattr(engine_args, "inline_diffusion", False) - ) - custom_pipeline_args = ( - engine_args.get("custom_pipeline_args") - if hasattr(engine_args, "get") - else getattr(engine_args, "custom_pipeline_args", None) - ) + if isinstance(plan.stage_cfg, BaseVllmOmniStageConfig): + inline_diffusion = plan.stage_cfg.stage_pipeline_config.inline_diffusion + custom_pipeline_args = getattr( + getattr(plan.stage_cfg, "diffusion_config", None), + "custom_pipeline_args", + None, + ) + else: + engine_args = getattr(plan.stage_cfg, "engine_args", {}) + inline_diffusion = ( + engine_args.get("inline_diffusion", False) + if hasattr(engine_args, "get") + else getattr(engine_args, "inline_diffusion", False) + ) + custom_pipeline_args = ( + engine_args.get("custom_pipeline_args") + if hasattr(engine_args, "get") + else getattr(engine_args, "custom_pipeline_args", None) + ) client, resources = launch_diffusion_stage_replica( model=self._model, stage_config=plan.stage_cfg, @@ -1028,7 +1053,7 @@ class DistStageRuntime(StageRuntime): for idx, stage_cfg in enumerate(self._stage_configs): stage_id = int(getattr(stage_cfg, "stage_id", idx)) - runtime_cfg = getattr(stage_cfg, "runtime", None) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None)) if runtime_cfg is None: continue if stage_id == target_stage_id: @@ -1100,13 +1125,9 @@ class DistStageRuntime(StageRuntime): if registered_stage_cfg is None: raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") - # Remote diffusion registration still transports the legacy mapping - # shape. Reconstruct and project that shape until its RFC #4021 cutover. - metadata = ( - extract_legacy_stage_metadata(OmegaConf.create(registered_stage_cfg)) - if plan.metadata.stage_type == "diffusion" - else copy.deepcopy(plan.metadata) - ) + # Registration is transport-only: the head-side typed plan remains + # authoritative for metadata and topology. + metadata = extract_stage_metadata_from_omni_stage_config(plan.stage_cfg) metadata.replica_id = plan.replica_id ctx = StageRemoteFactoryContext( stage_id=plan.metadata.stage_id, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index bc7f81132..a35f06202 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -1048,6 +1048,7 @@ def run_headless(args: TrackingNamespace) -> None: ) from vllm_omni.engine.stage_init_utils import ( build_engine_args_dict, + build_engine_args_dict_from_omni_stage_config, build_vllm_config, get_stage_connector_spec, inject_omni_kv_connector_config, @@ -1141,21 +1142,21 @@ def run_headless(args: TrackingNamespace) -> None: stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, - async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)), + async_chunk=bool( + getattr(getattr(stage_cfg, "connector_config", None), "async_chunk", None) + if hasattr(stage_cfg, "connector_config") + else stage_cfg.engine_args.get("async_chunk", False) + ), ) - # ``runtime_cfg`` is mostly inherited from the parent's - # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally - # bracket each replica's spawn below with setup_stage_devices so they - # don't all stack on cuda:0 (see ``per_replica_devices`` above). - # Headless startup still supplies the legacy OmegaConf stage shape through - # the stable adapter entry point. The implementation switches only when - # RFC #4021 threads structured stage configs through the launch plan. - engine_args_dict = build_engine_args_dict( - stage_cfg, - model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=getattr(args, "tokenizer", None), + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None) + ) + if hasattr(stage_cfg, "connector_config") + else build_engine_args_dict( + stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None) + ) ) inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 484d5d170..5b3f2d54b 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -248,10 +248,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return obj.get(key, default) return getattr(obj, key, default) + @classmethod + def _stage_model_metadata(cls, stage: Any) -> tuple[str | None, str | None]: + """Return ``(model_arch, model_stage)`` for typed or legacy stages.""" + model_config = cls._stage_get(stage, "model_config") + model_arch = cls._stage_get(model_config, "model_arch") + model_stage = cls._stage_get(stage, "model_stage") + + topology = cls._stage_get(stage, "stage_pipeline_config") + if model_stage is None: + model_stage = cls._stage_get(topology, "model_stage") + + engine_args = cls._stage_get(stage, "engine_args") + if model_arch is None: + model_arch = cls._stage_get(engine_args, "model_arch") + if model_stage is None: + model_stage = cls._stage_get(engine_args, "model_stage") + return model_arch, model_stage + def _has_minicpmo45_stage(self) -> bool: for stage in getattr(self.engine_client, "stage_configs", []) or []: - engine_args = self._stage_get(stage, "engine_args") - model_arch = self._stage_get(engine_args, "model_arch") + model_arch, _ = self._stage_model_metadata(stage) if model_arch == "MiniCPMO45OmniForConditionalGeneration": return True return False @@ -324,11 +341,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if not hasattr(sp, "output_kind"): continue - engine_args = self._stage_get(stage, "engine_args") - if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration": + model_arch, model_stage = self._stage_model_metadata(stage) + if model_arch != "MiniCPMO45OmniForConditionalGeneration": continue - model_stage = self._stage_get(engine_args, "model_stage") if model_stage == "llm": sp.output_kind = RequestOutputKind.FINAL_ONLY elif model_stage == "tts": diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 80f357651..8aa32899a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -36,6 +36,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler from vllm_omni.entrypoints.openai.protocol.audio import ( AudioResponse, @@ -76,6 +77,24 @@ _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS" _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS" _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS" + +def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]: + """Return speech routing metadata from typed or legacy stage configs.""" + topology = getattr(stage, "stage_pipeline_config", None) + if isinstance(topology, StagePipelineConfig): + return ( + getattr(topology, "model_stage", None), + getattr(getattr(stage, "model_config", None), "model_arch", None), + getattr(stage, "worker_type", None), + ) + engine_args = getattr(stage, "engine_args", None) + return ( + getattr(engine_args, "model_stage", None), + getattr(engine_args, "model_arch", None), + getattr(engine_args, "worker_type", None), + ) + + # TTS Configuration # # The stage-key -> model-type mapping is NOT declared here: it is derived from @@ -472,12 +491,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """Find and return the TTS stage config, or None if not found.""" tts_stage_keys = all_tts_stage_keys() entry_stage_archs = tts_entry_stage_archs() - all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs) + all_stages = frozenset( + model_stage + for stage in self.engine_client.stage_configs + if (model_stage := _stage_speech_metadata(stage)[0]) is not None + ) for stage in self.engine_client.stage_configs: - engine_args = stage.engine_args - model_stage = engine_args.model_stage - model_arch = getattr(engine_args, "model_arch", None) - worker_type = getattr(engine_args, "worker_type", None) + model_stage, model_arch, worker_type = _stage_speech_metadata(stage) if model_stage in tts_stage_keys: # Owning the stage key is not always enough: a model may be # speech-capable only in some deployment topologies. Ask the @@ -501,10 +521,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """ if self._tts_stage is None: return None - return detect_tts_model_type( - getattr(self._tts_stage.engine_args, "model_stage", None), - getattr(self._tts_stage.engine_args, "model_arch", None), - ) + return detect_tts_model_type(*_stage_speech_metadata(self._tts_stage)[:2]) def _compute_max_instructions_length(self) -> int: """Compute max instructions length with precedence: CLI > stage config > default. @@ -518,7 +535,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # 2. Try to get from TTS stage config if self._tts_stage is not None: - tts_args = getattr(self._tts_stage, "tts_args", {}) + topology = getattr(self._tts_stage, "stage_pipeline_config", None) + tts_args = ( + getattr(topology, "extras", {}).get("tts_args", {}) + if isinstance(topology, StagePipelineConfig) + else getattr(self._tts_stage, "tts_args", {}) + ) if "max_instructions_length" in tts_args: return tts_args["max_instructions_length"] @@ -1896,9 +1918,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # chat-template markers (im_start_token_id 151644). A raw-text # prompt produces a 1-token placeholder that crashes the talker's # prefill/decode handoff. Reject early with an actionable message. - stage_names = { - getattr(getattr(s, "engine_args", None), "model_stage", None) for s in self.engine_client.stage_configs - } + stage_names = {_stage_speech_metadata(s)[0] for s in self.engine_client.stage_configs} if "talker" in stage_names: raise ValueError( "The /v1/audio/speech endpoint is only supported for " diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 9a24a1ef9..717496362 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -53,6 +53,30 @@ logger = init_logger(__name__) _VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8 + +def _config_value(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, Mapping): + return config.get(key, default) + return getattr(config, key, default) + + +def _stage_diffusion_model_class_name(stage_config: Any) -> str | None: + """Resolve a stage's diffusion class across typed and legacy shapes.""" + diffusion_config = _config_value(stage_config, "diffusion_config") + model_class_name = _config_value(diffusion_config, "model_class_name") + if model_class_name: + return str(model_class_name) + + model_config = _config_value(stage_config, "model_config") + model_arch = _config_value(model_config, "model_arch") or _config_value(stage_config, "model_arch") + if model_arch: + return str(model_arch) + + engine_args = _config_value(stage_config, "engine_args", {}) + model_class_name = _config_value(engine_args, "model_class_name") + return str(model_class_name) if model_class_name else None + + if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -195,18 +219,7 @@ class OmniOpenAIServingVideo: model_class_name = getattr(od_config, "model_class_name", None) model_archs = [model_class_name] for stage_config in self.stage_configs or (): - stage_get = ( - stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) - ) - engine_args = stage_get("engine_args") or {} - model_archs.extend( - ( - stage_get("model_arch"), - engine_args.get("model_class_name") - if isinstance(engine_args, Mapping) - else getattr(engine_args, "model_class_name", None), - ) - ) + model_archs.append(_stage_diffusion_model_class_name(stage_config)) metadata_capability = any( get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs ) @@ -223,18 +236,7 @@ class OmniOpenAIServingVideo: od_config = self._resolve_diffusion_od_config() model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)] for stage_config in self.stage_configs or (): - stage_get = ( - stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) - ) - engine_args = stage_get("engine_args") or {} - model_archs.extend( - ( - stage_get("model_arch"), - engine_args.get("model_class_name") - if isinstance(engine_args, Mapping) - else getattr(engine_args, "model_class_name", None), - ) - ) + model_archs.append(_stage_diffusion_model_class_name(stage_config)) supported: set[str] = set() for model_arch in model_archs: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py index 0e31b633f..3b9e4a9a5 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py @@ -42,7 +42,11 @@ class AudexAdapter(ARTTSAdapter): from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot - stage = getattr(getattr(self.ctx.server._tts_stage, "engine_args", None), "model_stage", None) + stage_config = self.ctx.server._tts_stage + topology = getattr(stage_config, "stage_pipeline_config", None) + stage = getattr(topology, "model_stage", None) + if stage is None: + stage = getattr(getattr(stage_config, "engine_args", None), "model_stage", None) profile, folder = ( ("full", "checkpoint_folder_full") if stage == "audex_omni" else ("tts", "checkpoint_folder_audiogen") ) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index ff43b27a5..cb760e865 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -39,6 +39,21 @@ DEFAULT_TTS_LANGUAGES = frozenset( ) +def resolve_stage_model_path(engine_client: Any) -> str | None: + """Prefer typed stage model overrides, then legacy and served-model values.""" + stages = getattr(engine_client, "stage_configs", ()) or () + for stage in stages: + model_path = getattr(getattr(stage, "model_config", None), "model", None) + if model_path: + return str(model_path) + for stage in stages: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + return str(model_path) + model_path = getattr(engine_client, "model", None) + return str(model_path) if model_path else None + + def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str: """Stable hash of the real Stage 0 conditioning for the prefix cache. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py index 207d8b47b..33f16cf0a 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py @@ -9,7 +9,12 @@ import numpy as np from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + resolve_stage_model_path, +) if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -81,17 +86,7 @@ class HiggsAudioV2Adapter(ARTTSAdapter): from transformers import AutoProcessor - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - # Fallback: the orchestrator stores the served model id on the engine - # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args - # may not surface ``model`` when the deploy yaml doesn't set it per - # stage (the CLI-passed model id is the single source of truth). - model_path = getattr(self.engine_client, "model", None) + model_path = resolve_stage_model_path(self.engine_client) if model_path is None: raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs") processor = AutoProcessor.from_pretrained(model_path) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index 1e64f7a63..d993a7395 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -16,6 +16,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( PreparedRequest, apply_max_new_tokens, conditioning_cache_salt, + resolve_stage_model_path, ) _REF_CODE_CACHE_MAX_ENTRIES = 256 @@ -170,13 +171,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter): HiggsAudioV3TokenizerAdapter, ) - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - model_path = getattr(self.engine_client, "model", None) + model_path = resolve_stage_model_path(self.engine_client) if model_path is None: raise RuntimeError("higgs_audio_v3 serving could not resolve model path") tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py index 4473fce07..0ffd1ee83 100644 --- a/vllm_omni/entrypoints/openai/video/generation/helpers.py +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -58,6 +58,7 @@ from vllm_omni.entrypoints.openai.serving_video import ( ReferenceAudio, ReferenceImage, ReferenceVideo, + _stage_diffusion_model_class_name, ) from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER from vllm_omni.entrypoints.openai.stores import VIDEO_STORE @@ -136,10 +137,6 @@ def _config_get(config: Any, key: str, default: Any = None) -> Any: return getattr(config, key, default) -def _stage_engine_args(stage_cfg: Any) -> Any: - return _config_get(stage_cfg, "engine_args", {}) or {} - - def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: if not stage_configs: return [] @@ -150,7 +147,7 @@ def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: for stage_cfg in stage_configs: if get_stage_type(stage_cfg) != "diffusion": continue - model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name") + model_class_name = _stage_diffusion_model_class_name(stage_cfg) if not model_class_name: continue model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 562ed70d4..11fde2bf8 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -103,6 +103,14 @@ class ServingRealtimeRobotOpenPI: for stage_config in getattr(engine_client, "stage_configs", []) or []: if getattr(stage_config, "stage_type", None) != "diffusion": continue + # Typed diffusion stages keep model-owned OpenPI handshake + # metadata in diffusion_config.model_config. The out-of-process + # head only has this stage view because full od_config lives in + # the worker. + diffusion_config = getattr(stage_config, "diffusion_config", None) + model_config = getattr(diffusion_config, "model_config", None) + if model_config is not None: + break engine_args = getattr(stage_config, "engine_args", None) model_config = getattr(engine_args, "model_config", None) if model_config is not None: diff --git a/vllm_omni/entrypoints/serve/profile/utils.py b/vllm_omni/entrypoints/serve/profile/utils.py index f444a3534..d91ef2cad 100644 --- a/vllm_omni/entrypoints/serve/profile/utils.py +++ b/vllm_omni/entrypoints/serve/profile/utils.py @@ -4,18 +4,20 @@ def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool: - """Check if any stage has profiler_config set in its engine_args.""" + """Check whether any typed or legacy stage enables a profiler.""" if not stage_configs: return False for stage in stage_configs: - engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) - if engine_args is None: - continue profiler_config = ( - engine_args.get("profiler_config") - if isinstance(engine_args, dict) - else getattr(engine_args, "profiler_config", None) + stage.get("profiler_config") if isinstance(stage, dict) else getattr(stage, "profiler_config", None) ) + if profiler_config is None: + engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) + profiler_config = ( + engine_args.get("profiler_config") + if isinstance(engine_args, dict) + else getattr(engine_args, "profiler_config", None) + ) if profiler_config is not None: profiler = ( profiler_config.get("profiler") diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 281cbc55c..8c2ba0747 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -1,7 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import logging import os +from dataclasses import asdict, is_dataclass from multiprocessing import shared_memory as _shm from typing import Any @@ -274,6 +278,8 @@ def _to_dict(x: Any) -> dict[str, Any]: try: if isinstance(x, dict): return dict(x) + if is_dataclass(x) and not isinstance(x, type): + return asdict(x) return _omega_to_dict(x) except Exception: try: diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 39131fd58..74234d8d6 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -18,6 +18,13 @@ logger = init_logger(__name__) def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None: """Inject connector configuration into stage engine arguments.""" + typed_connector_config = getattr(stage, "connector_config", None) + if typed_connector_config is not None: + omni_conf_dict = dict(typed_connector_config.omni_kv_config or {}) + omni_conf_dict.update(connector_config=omni_conn_cfg, omni_from_stage=omni_from, omni_to_stage=omni_to) + typed_connector_config.omni_kv_config = omni_conf_dict + return + # Prepare omni_kv_config dict omni_conf_dict = {} try: From a769d9357b7800b521b25bc66b63a14020e647c1 Mon Sep 17 00:00:00 2001 From: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:50:38 +0800 Subject: [PATCH 21/99] [Bugfix][NPU] Fix talker_mtp ACL graph capture missing force_uniform_decode (#7008) Signed-off-by: Wallbreazzz Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Co-authored-by: Canlin Guo Co-authored-by: Hongsheng Liu --- vllm_omni/platforms/npu/worker/npu_ar_model_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 7bc6c21c6..f78b777a5 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -177,6 +177,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32), max_num_scheduled_tokens=1, use_cascade_attn=False, + force_uniform_decode=True, ) n = batch_desc.num_tokens ids = self.talker_mtp_input_ids.gpu[:n] From 4200768f5923b13f7d2c54af7d687d479e929573 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 16 Sep 2026 22:29:31 -0700 Subject: [PATCH 22/99] [Frontend] Scale Omni serving across API processes (#6923) Signed-off-by: Sy03 <1370724210@qq.com> --- docs/cli/serve.md | 69 +++ docs/design/module/entrypoints.md | 10 + docs/design/module/stage_runtime.md | 20 + .../examples/online_serving/qwen3_omni.md | 44 +- .../examples/online_serving/text_to_speech.md | 20 + examples/online_serving/qwen3_omni/README.md | 13 +- .../online_serving/text_to_speech/README.md | 20 +- .../test_async_omni_engine_stage_init.py | 414 ++++++++++++++++- tests/engine/test_stage_device_layout.py | 39 +- .../openai_api/test_api_server_guards.py | 191 ++++++++ .../openai_api/test_omni_sleep_wakeup.py | 2 + .../openai_api/test_serving_speech.py | 1 + .../openai_api/test_video_server.py | 3 + tests/entrypoints/test_serve.py | 289 ++++++++++++ tests/entrypoints/test_utils.py | 18 + .../deploy/qwen3_tts_high_concurrency.yaml | 9 + vllm_omni/engine/omni_engine_base.py | 36 +- vllm_omni/engine/stage_engine_startup.py | 31 +- vllm_omni/engine/stage_init_utils.py | 13 +- vllm_omni/engine/stage_runtime.py | 429 +++++++++++++----- vllm_omni/entrypoints/cli/serve.py | 279 +++++++++++- vllm_omni/entrypoints/openai/api_server.py | 138 +++++- vllm_omni/entrypoints/utils.py | 62 ++- vllm_omni/utils/tracking_parser.py | 20 +- .../worker/gpu_generation_model_runner.py | 2 +- vllm_omni/worker/gpu_model_runner.py | 6 +- 26 files changed, 1974 insertions(+), 204 deletions(-) diff --git a/docs/cli/serve.md b/docs/cli/serve.md index cf3d068a6..f77503091 100644 --- a/docs/cli/serve.md +++ b/docs/cli/serve.md @@ -1,5 +1,67 @@ # vllm-omni serve +## Multiple API frontends + +For local EngineCore pipelines with one or more stages, `--api-server-count` starts +multiple API frontend processes that share one parent-owned set of stage +engines: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ + --api-server-count 2 \ + --port 8091 +``` + +This mode is intended for workloads whose request parsing, multimodal input +processing, or response encoding can bottleneck a single frontend. It +currently requires local multi-process EngineCore stages with one local +process group per replica. A stage may still use multiple `num_replicas`. +It cannot be combined with diffusion, headless or +remote stages, intra-stage data parallelism, Ray, fault tolerance, elastic +expert parallelism, sleep mode, or runtime LoRA updating. Runtime voice upload +and deletion are also disabled because those mutations are process-local; +built-in voices, inline reference audio, and voices restored at startup remain +available. The `/v1/omni/sleep` and `/v1/omni/wakeup` control routes return HTTP +409 for the same reason: their bookkeeping is process-local while stage engines +are shared. + +A single-stage model pipeline is supported through the normal launch command; +this is distinct from the distributed `--stage-id` / `--headless` launch mode. +Every stage in the pipeline must be an EngineCore stage. Mixed AR-to-DiT +pipelines, including MiniMax-H3, remain unsupported even when their first +stage is an EngineCore stage. Diffusion support requires per-frontend channels +and shared-engine lifecycle handling in the diffusion backend. + +### Workloads and performance validation + +Each API process owns its frontend processor, Orchestrator, request state, +and per-replica EngineCore channels. Requests return through the channel of +the originating API process; model weights and GPU stage engines are shared. +This distributes frontend and orchestration work across CPU processes. + +| Workload | Potential benefit | Limitation | +| ---------- | ------------------- | ------------ | +| Concurrent cold image/video inputs on supported EngineCore pipelines | Parallel frontend processing when one API process is CPU-limited | A single cold request and GPU computation may be unchanged | +| Many replicas producing frequent text/audio outputs | Less API/Orchestrator work and GIL contention per process | Routing remains serial within each Orchestrator | +| Warm inputs or low-concurrency, GPU-limited generation | Additional frontend capacity | More API processes may add overhead without improving throughput | + +Here, a cold input means media or preprocessing caches miss, not that model +weights must be loaded or GPU kernels compiled. Each frontend has its own +processor/cache state, so cache duplication, CPU thread contention, and memory +usage also matter. Multi-API serving does not itself change the output queue, +move response encoding to an executor, or change the model's numerical scheduler. + +Compare `--api-server-count 1`, `2`, and higher counts with the same hardware, +model/stage configuration, media, sampling parameters, output lengths, and +offered load. Keep cold-cache and warmed-cache runs separate and report repeated +measurements. Record CPU and request counts for every API worker, input and +output queue delay, GPU utilization, first-output latency, and end-to-end latency. +Confirm requests reach every worker: connection reuse can concentrate traffic +on one process. Queue-only or single-API experiments do not establish a +multi-API throughput gain. The multi-replica workload in +[#4680](https://github.com/vllm-project/vllm-omni/issues/4680) remains a relevant +validation target. + ## Stage-based CLI quickstart The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process @@ -61,3 +123,10 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ ## Arguments --8<-- "docs/generated/argparse_omni/omni_serve.inc.md" + +Async video jobs use process-local metadata and task handles. With multiple API +workers, creating, listing, retrieving, downloading, and deleting these jobs +returns HTTP 409. This restriction is independent of the current diffusion +stage launch restriction: supporting multi-worker diffusion also requires a +shared job lifecycle. Missing frontend topology state returns HTTP 503 for +operations that require process-local state protection. diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index db62f9cb6..f37ddc6ba 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -94,6 +94,13 @@ This document owns offline API semantics, CLI and serve composition, supported OpenAI-compatible routes, request validation and normalization, response conversion, streaming/session behavior, and engine handoff. +When `--api-server-count` is greater than one, the serve composition root +starts the frontend process manager while a parent-owned `StageRuntime` +launches the shared local stage engines. Each frontend receives only its own +stage-client channel configuration and does not launch or retire backend +processes. CLI validation rejects unsupported distributed, diffusion, +fault-tolerant, elastic-EP, Ray, and runtime-LoRA combinations before startup. + It does not own configuration precedence, cross-stage routing, stage placement, payload implementation, or semantic error classification. `entrypoints/openai/errors.py` is an explicit primary-path exception owned by @@ -141,6 +148,9 @@ stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])` Streaming input pumps take an admission slot immediately before each EngineCore ADD or update, not while waiting for the next client chunk. Frontend abort keeps `request_states` until `generate()` consumes the terminal output. +Multi-API serving must preserve one client rank and one channel set per +frontend, include every frontend in readiness/failure observation, and keep +stage lifecycle ownership in the parent composition root. ## Promotion gate diff --git a/docs/design/module/stage_runtime.md b/docs/design/module/stage_runtime.md index aefba5a40..38255880f 100644 --- a/docs/design/module/stage_runtime.md +++ b/docs/design/module/stage_runtime.md @@ -77,6 +77,22 @@ readiness, one pool per logical stage, replica identity and membership, selection and affinity, stage-client/process lifecycle, liveness, draining, and shutdown. +Local EngineCore initialization and multi-API serving use +`StageRuntime.launch_stage_engines(num_api_servers=1)` as their common launch +entry. It resolves plans when not supplied by the runtime's device-group +scheduler, coordinates device locks and environment overlays, starts engines, +and rolls back resources on startup or attachment failure. Single-API clients +attach locally after readiness and take ownership of the engine resources; +there is no additional API subprocess. Diffusion and remote attachment retain +their backend-specific initialization. + +For local multi-API serving, the parent runtime owns and launches each stage +engine once. It supplies a distinct input/output channel for every API +frontend and stage replica; frontend workers attach clients to those engines +without acquiring stage-process ownership. Readiness observation and shutdown +remain parent-owned. This topology currently excludes diffusion and +remote/headless stages, intra-stage data parallelism, and Ray backends. + It does not own cross-stage request policy, model scheduler policy, connector transfer semantics, configuration precedence, or public error rendering. @@ -114,6 +130,10 @@ affinity, membership changes, abort routing, process failure, draining, and repeated shutdown. Cross-stage policy changes belong in `engine_orchestration.md`. +For multi-API changes, additionally exercise per-frontend channel assignment, +shared-engine startup handshakes, frontend failure observation, and parent +cleanup without duplicate stage shutdown. + ## Promotion gate - Reconcile terminology and ownership after #5441 reaches a final state. diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 2c85c33fe..654802c5c 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -2,7 +2,6 @@ Source . - ## 🛠️ Installation Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md) @@ -15,6 +14,19 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` +For a local deployment with multiple API frontend processes sharing one set +of stage engines, add `--api-server-count`: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ + --api-server-count 2 --port 8091 +``` + +This mode supports local EngineCore stages with one local process group per +replica, including stages configured with multiple `num_replicas`. Headless, +remote, diffusion, intra-replica data-parallel, and Ray stage deployments +remain single-frontend. + The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies. The bundled Qwen3-Omni setup explicitly sets `moe_backend: triton` for the Thinker and Talker stages. This keeps @@ -27,6 +39,7 @@ on the Thinker stage. Keep Talker and Code2Wav at `pipeline_parallel_size: 1`. T default asynchronous scheduling and audio chunk streaming can remain enabled. To explicitly utilize a custom deployment YAML, specify the configuration path: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/deploy_config_file @@ -50,7 +63,7 @@ Adopt the stage-based CLI architecture to independently instantiate execution pr The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via `CUDA_VISIBLE_DEVICES`. -**1. Stage 0 (Thinker + API server)** +#### 1. Stage 0 (Thinker + API server) ```bash CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -60,7 +73,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**2. Stage 1 (Talker)** +#### 2. Stage 1 (Talker) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -70,7 +83,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**3. Stage 2 (Code2Wav)** +#### 3. Stage 2 (Code2Wav) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -95,6 +108,9 @@ To experiment with another supported MoE backend, set `moe_backend` in the deployment YAML for both the Thinker and Talker stages. For the stage-based CLI, pass `--moe-backend` directly on each stage command. +For other stage settings, pass the option on that stage command. For example, +to change its GPU memory utilization: + ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --stage-id 1 \ @@ -107,11 +123,12 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ ### Send Multi-modal Request Get into the example folder + ```bash cd examples/online_serving/qwen3_omni ``` -#### Send request via python +#### Send request via python ```bash python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost" @@ -126,7 +143,6 @@ The Python client supports the following command-line arguments: - `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type is `use_audio`, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3` - `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` - For example, to use a local video file with custom prompt: ```bash @@ -136,19 +152,20 @@ python openai_chat_completion_client_for_multimodal_generation.py \ --prompt "What are the main activities shown in this video?" ``` -#### Send request via curl +#### Send request via curl ```bash bash run_curl_multimodal_generation.sh use_image ``` ## Modality control + You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance. ### Supported modalities | Modalities | Output | -|------------|--------| +| ------------ | -------- | | `["text"]` | Text only | | `["audio"]` | Audio only | | `["text", "audio"]` | Text + Audio | @@ -226,7 +243,9 @@ print(response.choices[1].message.audio) # Audio response ``` ## Streaming Output + If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally. + ```bash python openai_chat_completion_client_for_multimodal_generation.py \ --query-type use_image \ @@ -259,12 +278,14 @@ The convenience script launches both the vLLM server and Gradio demo together: ``` This script will: + 1. Start the vLLM server in the background 2. Wait for the server to be ready 3. Launch the Gradio demo 4. Handle cleanup when you press Ctrl+C The script supports the following arguments: + - `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct) - `--server-port`: Port for vLLM server (default: 8091) - `--gradio-port`: Port for Gradio demo (default: 7861) @@ -275,18 +296,19 @@ The script supports the following arguments: #### Option 2: Manual Launch (Two-Step Process) -**Step 1: Launch the vLLM API server** +##### Step 1: Launch the vLLM API server ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` If you have a custom deploy config file: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file ``` -**Step 2: Run the Gradio demo** +##### Step 2: Run the Gradio demo In a separate terminal: @@ -299,7 +321,7 @@ Then open `http://localhost:7861/` on your local browser to interact with the we The gradio script supports the following arguments: - `--model`: Model name/path (should match the server model) -- `--api-base`: Base URL for the vLLM API server (default: http://localhost:8091/v1) +- `--api-base`: Base URL for the vLLM API server (default: ) - `--ip`: Host/IP for Gradio server (default: 127.0.0.1) - `--port`: Port for Gradio server (default: 7861) - `--share`: Share the Gradio demo publicly (creates a public link) diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index e9499551f..c1107ad60 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -413,6 +413,22 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 ./qwen3_tts/run_server.sh Base ``` +For a local deployment with multiple API frontend processes sharing one set +of TTS stage engines, add `--api-server-count`: + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --omni --api-server-count 2 --port 8091 +``` + +This mode supports local EngineCore stages with one local process group per +replica, including stages configured with multiple `num_replicas`. Headless, +remote, intra-replica data-parallel, and Ray stage deployments are not +supported. Runtime voice upload and deletion are disabled with multiple +API frontends because their registries are process-local. Built-in voices, +inline `ref_audio`, and voices restored from `custom_voice_dir` at startup are +supported. + ### Executor backend Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override. @@ -478,6 +494,10 @@ checkpoint. When a request names an uploaded voice, the server infers `task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain CustomVoice speakers and require a CustomVoice checkpoint. +The runtime upload and delete routes require a single API frontend; with +`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices +from `custom_voice_dir` at startup. + ### Precomputed custom voices For reused Base voice-cloning speakers, precompute the reference artifacts once diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 107e1f098..8d6014db2 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -14,6 +14,17 @@ Please refer to [README.md](../../../README.md) vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` +For a local deployment with multiple API frontend processes sharing one set +of stage engines, add `--api-server-count`: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --api-server-count 2 --port 8091 +``` + +This mode currently supports local EngineCore stages. It is not available +with headless or remote stage deployments, and pure diffusion stages remain +single-frontend. + The default deployment configuration, situated at `vllm_omni/deploy/qwen3_omni_moe.yaml`, is resolved and loaded automatically via the model registry, obviating the `--deploy-config` flag in standard deployment topologies. Asynchronous chunk streaming operates as **enabled by default** within this bundled configuration. @@ -156,7 +167,7 @@ parser defaults). If you don't pass a flag, the YAML value wins. > chunked vs end-to-end modes (e.g. qwen3_tts code2wav) dispatch > automatically based on that bool — no extra flag or variant yaml is > needed. - +> > ⚠️ **For multi-stage models that share GPUs (qwen3_omni_moe by default > shares cuda:1 between stages 1 and 2), avoid using global memory flags.** > A global `--gpu-memory-utilization 0.85` would apply to every stage and diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index b749731e0..77e9ec4f1 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see ## Supported Models | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | -|---|---|---|---|---|---| +| --- | --- | --- | --- | --- | --- | | Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | @@ -781,6 +781,20 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 ./qwen3_tts/run_server.sh Base ``` +For a local deployment with multiple API frontend processes sharing one set +of TTS stage engines, add `--api-server-count`: + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --omni --api-server-count 2 --port 8091 +``` + +This mode supports local EngineCore stages only; headless or remote stage +deployments are not supported. Runtime voice upload and deletion are disabled +with multiple API frontends because their registries are process-local. +Built-in voices, inline `ref_audio`, and voices restored from `custom_voice_dir` +at startup are supported. + ### Executor backend Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override. @@ -835,6 +849,10 @@ checkpoint. When a request names an uploaded voice, the server infers `task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain CustomVoice speakers and require a CustomVoice checkpoint. +The runtime upload and delete routes require a single API frontend; with +`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices +from `custom_voice_dir` at startup. + ### Precomputed custom voices For reused Base voice-cloning speakers, precompute the reference artifacts once diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 637835122..75301ee0a 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -8,14 +8,17 @@ import json import os import time import types +from dataclasses import dataclass, field import pytest from omegaconf import OmegaConf +from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.config.omni_config import OmniStageRuntimeConfig from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_engine_startup import StageReplicaResources from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, ReplicaInitPlan, @@ -29,6 +32,32 @@ from vllm_omni.engine.stage_runtime import StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@dataclass +class _FakeParallelConfig: + enable_fault_tolerance: bool = False + enable_elastic_ep: bool = False + data_parallel_size: int = 1 + use_ray: bool = False + + +@dataclass +class _FakeVllmConfig: + parallel_config: _FakeParallelConfig = field(default_factory=_FakeParallelConfig) + + +@dataclass +class _FakeRuntimeConfig: + devices: str + + +@dataclass +class _FakeStageConfig: + stage_id: int + stage_type: str + engine_args: dict[str, object] + runtime: _FakeRuntimeConfig + + def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch): env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV" monkeypatch.delenv(env_key, raising=False) @@ -200,7 +229,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): ) -def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch): +def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch): import vllm_omni.engine.omni_engine_base as engine_mod engine = object.__new__(AsyncOmniEngine) @@ -218,6 +247,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp engine._omni_lb_policy = "random" engine.request_queue = types.SimpleNamespace() engine._log_stats = True + engine._client_config = engine_mod.OmniClientConfig(client_count=2, client_index=1, stage_addresses={}) engine._parallel_stage_init = False captured: dict[str, object] = {} @@ -233,6 +263,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp assert captured["stage_init_timeout"] == 7 assert captured["log_stats"] is True + assert captured["client_config"] is engine._client_config def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch): @@ -833,6 +864,289 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch): assert captured["client_log_stats"] is True +def test_stage_runtime_attaches_external_llm_client_without_launching_engine(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + client_addresses = { + "input_address": "ipc://stage0-input-1", + "output_address": "ipc://stage0-output-1", + } + client_config = { + "client_count": 2, + "client_index": 1, + "stage_addresses": {0: {0: client_addresses}}, + } + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + client_config=client_config, + ) + parallel_config = _FakeParallelConfig() + vllm_config = _FakeVllmConfig(parallel_config=parallel_config) + plan = _make_llm_plan(0, stage_id=0, vllm_config=vllm_config).replicas[0] + captured: dict[str, object] = {} + stage_client = object() + + def _capture_client(**kwargs): + captured.update(kwargs) + return stage_client + + monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _capture_client) + monkeypatch.setattr( + runtime_mod, + "launch_stage_replica", + lambda **_kwargs: pytest.fail("external client must not launch a stage engine"), + ) + + assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client + assert captured["client_addresses"] == client_addresses + assert captured["client_count"] == 2 + assert captured["client_index"] == 1 + assert not hasattr(parallel_config, "_api_process_count") + assert not hasattr(parallel_config, "_api_process_rank") + + +@pytest.mark.parametrize("stage_ids", [(0,), (0, 1)], ids=["single-stage", "multi-stage"]) +@pytest.mark.parametrize("client_count", [1, 2]) +def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeypatch, stage_ids, client_count): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + stage_plans = [] + for stage_id in stage_ids: + parallel_config = _FakeParallelConfig() + plan = _make_llm_plan( + stage_id, + stage_id=stage_id, + vllm_config=_FakeVllmConfig(parallel_config=parallel_config), + ) + plan.replicas[0].engine_args_dict = {} + stage_plans.append(plan) + stage_plans[0].replicas[0].metadata.runtime_cfg = {"env": {"VLLM_OMNI_TEST_STAGE_RUNTIME_ENV": "enabled"}} + + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id)) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + + events: list[tuple[str, int]] = [] + captured_launch_kwargs: list[dict[str, object]] = [] + captured_launch_env: list[str | None] = [] + + @contextlib.contextmanager + def _fake_launch_stage_replica(**kwargs): + stage_id = kwargs["stage_id"] + events.append(("enter", stage_id)) + captured_launch_kwargs.append(kwargs) + captured_launch_env.append(os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV")) + yield StageReplicaResources( + addresses=EngineZmqAddresses( + inputs=[f"ipc://stage{stage_id}-input-{idx}" for idx in range(client_count)], + outputs=[f"ipc://stage{stage_id}-output-{idx}" for idx in range(client_count)], + ), + ) + events.append(("exit", stage_id)) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica) + + with runtime.launch_stage_engines(client_count) as launch: + expected = ( + [("enter", stage_id) for stage_id in stage_ids] + if client_count > 1 + else [(event, stage_id) for stage_id in stage_ids for event in ("enter", "exit")] + ) + assert events == expected + for client_index, config in enumerate(launch.client_configs): + assert config["client_count"] == client_count + assert config["client_index"] == client_index + for stage_id in stage_ids: + assert config["stage_addresses"][stage_id][0] == { + "input_address": f"ipc://stage{stage_id}-input-{client_index}", + "output_address": f"ipc://stage{stage_id}-output-{client_index}", + } + + assert events == ( + [(event, stage_id) for event in ("enter", "exit") for stage_id in stage_ids] if client_count > 1 else expected + ) + assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_count") + assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_rank") + assert all( + kwargs["watched_frontend_processes"] is (launch.watched_frontend_processes if client_count > 1 else None) + for kwargs in captured_launch_kwargs + ) + assert captured_launch_env == ["enabled" if stage_id == 0 else None for stage_id in stage_ids] + assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None + + +@pytest.mark.parametrize("client_count", [1, 2]) +def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count): + import fcntl + + import vllm_omni.engine.stage_init_utils as init_utils + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)] + for stage_id, plan in enumerate(plans): + plan.replicas[0].engine_args_dict = {"tensor_parallel_size": stage_id + 1} + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda sid, _cfg: "0" if sid == 0 else "0,1") + monkeypatch.setattr( + init_utils, "device_init_lock_path", lambda device_id: str(tmp_path / f"device-{device_id}.lock") + ) + monkeypatch.setattr( + init_utils.time, "sleep", lambda _: pytest.fail("Overlapping stages must not wait on their own lock") + ) + + @contextlib.contextmanager + def launch(**kwargs): + # Real flock on a separate descriptor proves the runtime holds each lock. + for device_id in range(kwargs["stage_id"] + 1): + with open(tmp_path / f"device-{device_id}.lock", "a") as probe: + with pytest.raises(BlockingIOError): + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + yield StageReplicaResources( + addresses=EngineZmqAddresses( + inputs=[f"ipc://input-{i}" for i in range(client_count)], + outputs=[f"ipc://output-{i}" for i in range(client_count)], + ) + ) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch) + with runtime.launch_stage_engines(client_count): + pass + for device_id in (0, 1): + with open(tmp_path / f"device-{device_id}.lock", "a") as probe: + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(probe, fcntl.LOCK_UN) + + +@pytest.mark.parametrize("diffusion_stage_id", [0, 1], ids=["diffusion-only", "enginecore-to-diffusion"]) +def test_stage_runtime_multi_api_rejects_diffusion_before_launch(monkeypatch, diffusion_stage_id): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + stage_plans = [ + _make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) + for stage_id in range(diffusion_stage_id) + ] + stage_plans.append(_make_diffusion_plan(diffusion_stage_id, stage_id=diffusion_stage_id)) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) + monkeypatch.setattr( + runtime_mod, + "acquire_device_locks", + lambda *_args: pytest.fail("unsupported topology must fail before acquiring devices"), + ) + monkeypatch.setattr( + runtime_mod, + "launch_stage_replica", + lambda **_kwargs: pytest.fail("unsupported topology must not partially launch EngineCore stages"), + ) + + with pytest.raises(ValueError, match="diffusion stage\\(s\\) are not supported"): + with runtime.launch_stage_engines(2): + pytest.fail("unsupported topology must not yield a launch handle") + + +def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + stage_plan = _make_llm_plan( + 0, + stage_id=0, + vllm_config=_FakeVllmConfig(), + ) + replica = stage_plan.replicas[0] + replica.engine_args_dict = {} + replica.metadata.runtime_cfg = {"devices": "0"} + + device_env = runtime_mod.current_omni_platform.device_control_env_var + monkeypatch.setenv(device_env, "5") + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan]) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + captured_devices: list[str | None] = [] + + @contextlib.contextmanager + def _fake_launch_stage_replica(**kwargs): + captured_devices.append(kwargs["stage_visible_devices"]) + yield StageReplicaResources( + addresses=EngineZmqAddresses( + inputs=["ipc://input-0", "ipc://input-1"], + outputs=["ipc://output-0", "ipc://output-1"], + ), + ) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica) + + with runtime.launch_stage_engines(2): + pass + + assert captured_devices == ["5"] + assert os.environ[device_env] == "5" + + +def test_stage_runtime_multi_api_rejects_elastic_ep(monkeypatch): + runtime = _make_stage_runtime() + stage_plan = _make_llm_plan( + 0, + stage_id=0, + vllm_config=_FakeVllmConfig(parallel_config=_FakeParallelConfig(enable_elastic_ep=True)), + ) + stage_plan.replicas[0].engine_args_dict = {} + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan]) + + with pytest.raises(ValueError, match="enable-elastic-ep"): + with runtime.launch_stage_engines(2): + pass + + +def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_exit(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()) + stage_plan.replicas[0].engine_args_dict = {} + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan]) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + + events: list[str] = [] + + class _Manager: + def shutdown(self) -> None: + events.append("shutdown") + + @contextlib.contextmanager + def _fake_launch_stage_replica(**_kwargs): + try: + yield StageReplicaResources( + manager=_Manager(), + addresses=EngineZmqAddresses( + inputs=["tcp://127.0.0.1:0", "tcp://127.0.0.1:1"], + outputs=["tcp://127.0.0.1:2", "tcp://127.0.0.1:3"], + ), + ) + except RuntimeError as exc: + events.append(f"exceptional-exit:{exc}") + raise + else: + events.append("normal-exit") + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica) + + with pytest.raises(RuntimeError, match="deferred TCP addresses"): + with runtime.launch_stage_engines(2): + pass + + assert events == [ + "shutdown", + "exceptional-exit:Stage 0 returned deferred TCP addresses; multi-API launch requires fixed ports or IPC addresses", + ] + + def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod @@ -1079,7 +1393,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag num_replicas=1, launch_mode="local", stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")), - metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}), + metadata=types.SimpleNamespace(stage_id=0, stage_type="llm", runtime_cfg={"devices": "0"}), stage_connector_spec={}, omni_kv_connector=(None, None, None), stage_vllm_config=fake_vllm_config, @@ -1933,6 +2247,102 @@ def test_port_from_zmq_address_parsing(): assert _port_from_zmq_address("tcp://host:not-a-port") is None +@pytest.mark.parametrize("parallel", [False, True]) +@pytest.mark.parametrize("failure", ["ready", "attach", None]) +def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, failure): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + runtime._parallel_stage_init = parallel + plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0] + plan.engine_args_dict = {} + events = [] + + class Manager: + def shutdown(self): + events.append("shutdown") + + manager = Manager() + addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"]) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0") + + def acquire(*_): + events.append("lock") + return [42] + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire) + monkeypatch.setattr(runtime_mod, "release_device_locks", lambda _: events.append("unlock")) + + @contextlib.contextmanager + def launch(**kwargs): + assert kwargs["num_api_servers"] == 1 + assert kwargs["omni_parallel_stage_init"] is parallel + events.append("spawn") + yield StageReplicaResources(manager=manager, addresses=addresses) + assert runtime._replica_launch_lock.locked() is False + events.append("ready") + if failure == "ready": + raise RuntimeError("ready") + + def attach(**kwargs): + assert "ready" in events + assert kwargs["engine_manager"] is manager + events.append("attach") + if failure == "attach": + raise RuntimeError("attach") + return manager + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch) + monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", attach) + if failure: + with pytest.raises(RuntimeError, match=failure): + runtime._initialize_local_llm_replica(plan, 1) + assert events.count("shutdown") == 1 + else: + assert runtime._initialize_local_llm_replica(plan, 1) is manager + assert "shutdown" not in events + assert events.count("lock") == events.count("unlock") == (0 if parallel else 1) + + +@pytest.mark.parametrize("client_count", [1, 2]) +def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + runtime._parallel_stage_init = True + plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()) + plan.replicas[0].engine_args_dict = {} + events = [] + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [plan]) + monkeypatch.setattr(runtime, "_reject_unguardable_executors", lambda _: events.append("guard")) + monkeypatch.setattr(runtime, "_run_stage_admission", lambda _: events.append("admit")) + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0") + + def forbidden_lock(*_): + pytest.fail("parallel initialization must use child phase locks") + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", forbidden_lock) + + @contextlib.contextmanager + def launch(**kwargs): + assert events == ["guard", "admit"] + assert kwargs["omni_parallel_stage_init"] is True + events.append("spawn") + yield StageReplicaResources( + addresses=EngineZmqAddresses( + inputs=[f"ipc://in{i}" for i in range(client_count)], + outputs=[f"ipc://out{i}" for i in range(client_count)], + ) + ) + assert not runtime._replica_launch_lock.locked() + events.append("ready") + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch) + with runtime.launch_stage_engines(client_count): + pass + assert events == ["guard", "admit", "spawn", "ready"] + + @pytest.mark.parametrize( "stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)] ) diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py index beb66dfba..6dd30373d 100644 --- a/tests/engine/test_stage_device_layout.py +++ b/tests/engine/test_stage_device_layout.py @@ -18,6 +18,7 @@ import types from unittest import mock import pytest +from pytest_mock import MockerFixture from vllm_omni.engine import stage_init_utils from vllm_omni.engine.stage_init_utils import ( @@ -240,7 +241,7 @@ def test_build_vllm_config_fails_before_engine_config_on_mismatch(): get_class.assert_not_called() -def test_build_vllm_config_proceeds_on_consistent_layout(): +def test_build_vllm_config_proceeds_on_consistent_layout(mocker: MockerFixture) -> None: """The guard must not false-positive: a consistent single-GPU stage flows past it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual.""" stage = _stage(1, devices="1") @@ -249,18 +250,30 @@ def test_build_vllm_config_proceeds_on_consistent_layout(): model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()), ) sentinel_executor = object() - with ( - mock.patch.object( - stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config - ) as create_engine_config, - mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor), - mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant), - ): - vllm_config, executor_class = build_vllm_config( - stage, - model="dummy-model", - engine_args_dict={"tensor_parallel_size": 1}, - ) + captured_api_process_config: dict[str, int] = {} + + def capture_create_engine_config(engine_args, *args, **kwargs): + captured_api_process_config["count"] = engine_args._api_process_count + captured_api_process_config["rank"] = engine_args._api_process_rank + return fake_config + + create_engine_config = mocker.patch.object( + stage_init_utils.OmniEngineArgs, + "create_engine_config", + autospec=True, + side_effect=capture_create_engine_config, + ) + mocker.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor) + mocker.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant) + + vllm_config, executor_class = build_vllm_config( + stage, + model="dummy-model", + engine_args_dict={"tensor_parallel_size": 1}, + api_process_count=3, + api_process_rank=2, + ) create_engine_config.assert_called_once() assert vllm_config is fake_config assert executor_class is sentinel_executor + assert captured_api_process_config == {"count": 3, "rank": 2} diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 6e28bb5b1..ebc08662f 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -42,7 +42,10 @@ from __future__ import annotations import asyncio import json +import signal +from argparse import Namespace from contextlib import asynccontextmanager +from dataclasses import dataclass from types import SimpleNamespace import pytest @@ -59,6 +62,66 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +@dataclass +class _FakeParallelConfig: + _api_process_rank: int + + +@dataclass +class _FakeVllmConfig: + parallel_config: _FakeParallelConfig + + +def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker): + args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}]) + death_signal = mocker.patch.object(api_server, "set_death_signal") + + with pytest.raises(RuntimeError, match="client configuration"): + api_server.run_omni_api_server_worker_proc( + "127.0.0.1:8000", + mocker.Mock(), + args, + {"client_count": 2, "client_index": -1}, + ) + + death_signal.assert_called_once_with(signal.SIGTERM) + + +@pytest.mark.asyncio +async def test_omni_run_server_worker_forwards_client_config(monkeypatch): + captured: dict[str, object] = {} + + @asynccontextmanager + async def fake_build_async_omni(*_args, **kwargs): + captured.update(kwargs) + yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1))) + + monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni) + monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI()) + monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0)) + monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None) + + async def fake_storage_start() -> None: + return None + + monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start) + + async def fake_serve_http(app, **_kwargs): + assert app.state.api_server_count == 2 + task = asyncio.create_task(asyncio.sleep(0)) + await asyncio.sleep(0) + return task + + monkeypatch.setattr(api_server, "serve_http", fake_serve_http) + + args = _minimal_args(model="dummy") + config = {"client_count": 2, "client_index": 1, "stage_addresses": {}} + + await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config) + assert captured["client_config"] is config + + # FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here # would invent ("HEAD", path) entries for every GET and fail the current census. # TODO: extend this set and the expected census if @router.head / @router.options appear. @@ -626,6 +689,48 @@ async def test_realtime_route_defaults_to_configured_duplex_handler( assert calls == [expected_handler] +@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime", "/v1/realtime?duplex=1"]) +def test_multi_api_duplex_reconnect_rejected_on_each_worker(path: str) -> None: + class SessionHandler: + async def handle_realtime_session(self, _websocket) -> None: + pytest.fail("Multi-API requests must not access process-local sessions") + + class PendingWarmup: + def is_set(self) -> bool: + raise AssertionError("Multi-API rejection must not wait for duplex warmup") + + # Separate apps represent workers with independent frontend registries. + # Opening on A then reconnecting to B must reject at the same boundary. + for worker_index in (0, 1): + app = FastAPI() + app.state.api_server_count = 2 + app.state.api_server_index = worker_index + app.state.openai_serving_duplex = SessionHandler() + app.state.duplex_warmup_done = PendingWarmup() + app.include_router(api_server.router) + with TestClient(app) as client, client.websocket_connect(path) as websocket: + assert websocket.receive_json()["code"] == "multi_api_duplex_unsupported" + with pytest.raises(WebSocketDisconnect) as exc: + websocket.receive_text() + assert exc.value.code == 1008 + + +@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"]) +def test_single_api_duplex_session_still_reaches_handler(path: str) -> None: + class SessionHandler: + async def handle_realtime_session(self, websocket) -> None: + await websocket.accept() + await websocket.send_json({"type": "session.created"}) + await websocket.close() + + app = FastAPI() + app.state.api_server_count = 1 + app.state.openai_serving_duplex = SessionHandler() + app.include_router(api_server.router) + with TestClient(app) as client, client.websocket_connect(path) as websocket: + assert websocket.receive_json() == {"type": "session.created"} + + def test_health_without_engine_returns_stable_unhealthy_response() -> None: """Lock ``/health`` with no engine initialized. @@ -715,6 +820,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None: assert exc_info.value.detail == "The model does not support Speech API" +@pytest.mark.asyncio +async def test_multi_api_rejects_runtime_voice_upload() -> None: + app = FastAPI() + app.state.api_server_count = 2 + raw_request = _request_for(app, method="POST", path="/v1/audio/voices") + + with pytest.raises(HTTPException) as exc_info: + await api_server.upload_voice( + raw_request, + audio_sample=None, + speaker_embedding=None, + consent="consent-id", + name="probe", + ) + + assert exc_info.value.status_code == 409 + assert "process-local frontend state" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_multi_api_rejects_sleep_route() -> None: + app = FastAPI() + app.state.api_server_count = 2 + raw_request = _request_for(app, method="POST", path="/v1/omni/sleep") + + with pytest.raises(HTTPException) as exc_info: + await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request) + + assert exc_info.value.status_code == 409 + assert "process-local frontend state" in exc_info.value.detail + + def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None: """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers. @@ -937,3 +1074,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: must_be_wired=_MULTISTAGE_MUST_BE_WIRED, must_be_none=_MULTISTAGE_MUST_BE_NONE, ) + + +@pytest.mark.parametrize("count", [None, 0, -1, "2", True]) +def test_process_local_guard_rejects_unknown_topology(count): + app = FastAPI() + if count is not None: + app.state.api_server_count = count + # CLI args cannot substitute for missing/invalid actual runtime state. + app.state.args = Namespace(api_server_count=1) + with pytest.raises(HTTPException) as exc: + api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep") + assert exc.value.status_code == 503 + + +@pytest.mark.parametrize( + "requested,config,expected", + [(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)], +) +def test_api_topology_resolved_from_launch(requested, config, expected): + assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected + + +@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})]) +def test_api_topology_rejects_inconsistent_launch(requested, config): + with pytest.raises(ValueError): + api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) + + +@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)]) +@pytest.mark.parametrize( + "method,path", + [ + ("POST", "/v1/videos"), + ("GET", "/v1/videos"), + ("GET", "/v1/videos/job"), + ("GET", "/v1/videos/job/content"), + ("DELETE", "/v1/videos/job"), + ], +) +def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path): + from fastapi.testclient import TestClient + + app = FastAPI() + if count is not None: + app.state.api_server_count = count + app.include_router(api_server.router) + + async def forbidden_form(): + pytest.fail("video form processing must not run before topology rejection") + + app.dependency_overrides[api_server._parse_video_form] = forbidden_form + with TestClient(app) as client: + response = client.request(method, path) + assert response.status_code == status diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index 615248fcf..ba5925d8f 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py @@ -20,6 +20,7 @@ def _make_app(engine_client): from vllm_omni.entrypoints.openai.api_server import router app = FastAPI() + app.state.api_server_count = 1 app.include_router(router) app.state.engine_client = engine_client app.state.sleeping_stages = set() @@ -253,6 +254,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker): ) app = FastAPI() + app.state.api_server_count = 1 app.include_router(router) mocker.patch( diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 5ac895d86..7cdeebc39 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3826,6 +3826,7 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request: app = FastAPI() app.state.openai_serving_speech = handler + app.state.api_server_count = 1 scope = { "type": "http", "app": app, diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 0b44cabdf..715c7ca04 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -463,6 +463,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch) @pytest.fixture def test_client(): app = FastAPI() + app.state.api_server_count = 1 app.include_router(router) app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion( diffusion_engine=FakeAsyncOmni(), @@ -1804,6 +1805,7 @@ def test_action_extraction_accepts_multimodal_actions_payload(): def test_missing_handler_returns_503(): app = FastAPI() + app.state.api_server_count = 1 app.include_router(router) app.state.openai_serving_video = None client = TestClient(app) @@ -2790,6 +2792,7 @@ def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, tes def test_sync_missing_handler_returns_503(): app = FastAPI() + app.state.api_server_count = 1 app.include_router(router) app.state.openai_serving_video = None client = TestClient(app) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 8d3766df0..7f8b6a4a9 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -6,14 +6,19 @@ from __future__ import annotations import argparse +from dataclasses import dataclass +from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from typing import Any import pytest from pytest_mock import MockerFixture +from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.engine.stage_engine_startup import StageReplicaResources +from vllm_omni.engine.stage_runtime import StageEngineLaunch from vllm_omni.entrypoints.cli.serve import ( OmniServeCommand, _parse_stage_overrides, @@ -25,6 +30,31 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@dataclass(frozen=True) +class _FakeProcess: + sentinel: str + exitcode: int | None + name: str + pid: int + + +@dataclass +class _FakeProcessOwner: + processes: list[_FakeProcess] + + +@dataclass +class _FakeStageEngineArgs: + async_chunk: bool + enable_sleep_mode: bool + + +@dataclass +class _FakeStageConfig: + stage_id: int + engine_args: _FakeStageEngineArgs + + def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution: return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages)) @@ -44,6 +74,8 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: explicit = args.get_explicit_kwargs_dict() assert args.get_explicit_kwargs_dict() assert not explicit["async_chunk"] + assert (args.api_server_count or 1) == 1 + assert "api_server_count" not in explicit def _parse_serve_args(argv: list[str]) -> TrackingNamespace: @@ -169,6 +201,244 @@ def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value]) +def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2"]) + args._omni_stage_client_configs = [{"stage_addresses": {0: {0: {"input_address": "ipc://input"}}}}] + + restored = ForkingPickler.loads(ForkingPickler.dumps(args)) + + assert restored.api_server_count == 2 + assert restored.get_explicit_kwargs_dict()["api_server_count"] == 2 + assert restored._omni_stage_client_configs == args._omni_stage_client_configs + + +def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + cmd = OmniServeCommand() + cmd.subparser_init(subparsers) + args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"]) + + mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True) + validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args") + + with pytest.raises(ValueError, match="not supported for diffusion"): + cmd.validate(args) + + validate.assert_not_called() + + +def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: MockerFixture) -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + cmd = OmniServeCommand() + cmd.subparser_init(subparsers) + args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2", "--enable-sleep-mode"]) + + mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False) + mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args") + + with pytest.raises(ValueError, match="enable-sleep-mode"): + cmd.validate(args) + + +@pytest.mark.parametrize("typed", [False, True]) +@pytest.mark.parametrize("downstream_async_chunk", [False, True]) +def test_build_multi_api_stage_runtime_matches_current_constructor( + mocker: MockerFixture, typed: bool, downstream_async_chunk: bool +) -> None: + from vllm_omni.entrypoints.cli import serve as serve_module + + args = TrackingNamespace( + argparse.Namespace( + model="dummy-model", + model_tag=None, + stage_init_timeout=1, + tokenizer=None, + disable_log_stats=False, + ), + frozenset({"model"}), + ) + if typed: + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + + stages = [ + VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar")) + for i in range(2) + ] + stages[1].connector_config.async_chunk = downstream_async_chunk + else: + stages = [ + _FakeStageConfig( + stage_id=i, + engine_args=_FakeStageEngineArgs( + async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False + ), + ) + for i in range(2) + ] + mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model") + mocker.patch( + "vllm_omni.config.config_factory.with_trust_remote_code_override", + side_effect=lambda kwargs, _trust: kwargs, + ) + mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={}) + mocker.patch( + "vllm_omni.config.resolver.resolve_omni_config", + return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)), + ) + + runtime = serve_module._build_multi_api_stage_runtime(args, 2) + + assert runtime._client_count == 1 + assert runtime._client_index == 0 + assert runtime._async_chunk is downstream_async_chunk + + +@pytest.mark.parametrize("typed", [False, True]) +def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config( + mocker: MockerFixture, typed: bool +) -> None: + from vllm_omni.entrypoints.cli import serve as serve_module + + args = TrackingNamespace( + argparse.Namespace( + model="dummy-model", + model_tag=None, + stage_init_timeout=1, + tokenizer=None, + disable_log_stats=False, + ), + frozenset({"model"}), + ) + stage_config = _FakeStageConfig( + stage_id=3, + engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True), + ) + if typed: + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + + typed_stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=3, model_stage="ar")) + typed_stage.model_config.enable_sleep_mode = True + stage_config = typed_stage + mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model") + mocker.patch( + "vllm_omni.config.config_factory.with_trust_remote_code_override", + side_effect=lambda kwargs, _trust: kwargs, + ) + mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={}) + mocker.patch( + "vllm_omni.config.resolver.resolve_omni_config", + return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)), + ) + + with pytest.raises(ValueError, match=r"sleep mode.*stage\(s\) \[3\]"): + serve_module._build_multi_api_stage_runtime(args, 2) + + +def test_run_multi_api_server_omni_starts_workers_after_shared_engine_launch(mocker: MockerFixture) -> None: + from vllm_omni.entrypoints.cli import serve as serve_module + + args = TrackingNamespace( + argparse.Namespace(api_server_count=2, shutdown_timeout=1), + frozenset({"api_server_count", "shutdown_timeout"}), + ) + socket = mocker.Mock() + primary_addresses = EngineZmqAddresses( + inputs=["ipc://input-0", "ipc://input-1"], + outputs=["ipc://output-0", "ipc://output-1"], + ) + engine_launch = StageEngineLaunch( + client_configs=[{"client_count": 2, "client_index": index, "stage_addresses": {}} for index in range(2)], + resources=[StageReplicaResources(addresses=primary_addresses)], + ) + runtime = mocker.MagicMock() + runtime.launch_stage_engines.return_value.__enter__.return_value = engine_launch + runtime.launch_stage_engines.return_value.__exit__.return_value = False + + manager = mocker.Mock() + manager.processes = [mocker.Mock()] + manager.gather_actual_addresses.return_value = (primary_addresses.inputs, primary_addresses.outputs) + + mocker.patch.object(serve_module, "_build_multi_api_stage_runtime", return_value=runtime) + mocker.patch.object(serve_module, "_wait_for_multi_api_server_completion") + mocker.patch("signal.signal") + mocker.patch("vllm.entrypoints.openai.api_server.setup_server", return_value=("127.0.0.1:8000", socket)) + mocker.patch("vllm.v1.metrics.prometheus.setup_multiprocess_prometheus") + start_manager = mocker.patch.object(serve_module, "_start_api_server_process_manager", return_value=manager) + + serve_module.run_multi_api_server_omni(args) + + assert args._omni_stage_client_configs is engine_launch.client_configs + start_manager.assert_called_once() + manager_kwargs = start_manager.call_args.kwargs + assert manager_kwargs["cleanup_timeout"] == 1 + assert manager_kwargs["num_servers"] == 2 + assert manager_kwargs["input_addresses"] == ["ipc://input-0", "ipc://input-1"] + assert manager_kwargs["target_server_fn"] is serve_module.run_omni_api_server_worker_proc + assert engine_launch.watched_frontend_processes == manager.processes + runtime.shutdown.assert_called_once_with() + socket.close.assert_called_once_with() + + +def test_start_api_server_process_manager_cleans_up_partial_start(mocker: MockerFixture) -> None: + from vllm_omni.entrypoints.cli import serve as serve_module + + first_process = mocker.Mock(name="first_process") + second_process = mocker.Mock(name="second_process") + first_process.pid = 1234 + second_process.pid = None + second_process.start.side_effect = RuntimeError("spawn failed") + parent_pipes = [mocker.Mock(name="parent_pipe_0"), mocker.Mock(name="parent_pipe_1")] + child_pipes = [mocker.Mock(name="child_pipe_0"), mocker.Mock(name="child_pipe_1")] + spawn_context = mocker.Mock() + spawn_context.Process.side_effect = [first_process, second_process] + spawn_context.Pipe.side_effect = list(zip(parent_pipes, child_pipes)) + mocker.patch("vllm.v1.utils.multiprocessing.get_context", return_value=spawn_context) + shutdown = mocker.patch("vllm.v1.utils.shutdown") + + with pytest.raises(RuntimeError, match="spawn failed"): + serve_module._start_api_server_process_manager( + cleanup_timeout=1.5, + listen_address="127.0.0.1:8000", + sock=mocker.Mock(), + args=argparse.Namespace(), + num_servers=2, + input_addresses=["ipc://input-0", "ipc://input-1"], + output_addresses=["ipc://output-0", "ipc://output-1"], + target_server_fn=mocker.Mock(), + ) + + first_process.start.assert_called_once_with() + second_process.start.assert_called_once_with() + for pipe in parent_pipes: + pipe.close.assert_called_once_with() + shutdown.assert_called_once_with([first_process], timeout=1.5) + + +def test_wait_for_multi_api_server_completion_rejects_engine_failure( + mocker: MockerFixture, +) -> None: + from vllm_omni.entrypoints.cli import serve as serve_module + + api_process = _FakeProcess(sentinel="api-0", exitcode=None, name="api-0", pid=10) + engine_process = _FakeProcess(sentinel="engine-0", exitcode=4, name="engine-0", pid=20) + manager = _FakeProcessOwner(processes=[api_process]) + engine_launch = StageEngineLaunch( + client_configs=[], + resources=[StageReplicaResources(manager=_FakeProcessOwner(processes=[engine_process]))], + ) + mocker.patch("multiprocessing.connection.wait", return_value=["engine-0"]) + + with pytest.raises(RuntimeError, match=r"Shared stage engine process engine-0 .* exited with code 4"): + serve_module._wait_for_multi_api_server_completion(manager, engine_launch) + + def test_serve_parser_rejects_stage_configs_path() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") @@ -847,3 +1117,22 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur with pytest.raises(RuntimeError, match=r"exited with code 137"): run_headless(_make_headless_args(stage_id=1)) + + +@pytest.mark.parametrize("name", ["unfiltered_ns", "explicit_keys", "missing", "__dict__"]) +def test_tracking_namespace_uninitialized_access_raises_attribute_error(name): + namespace = object.__new__(TrackingNamespace) + with pytest.raises(AttributeError): + getattr(namespace, name) + + +@pytest.mark.parametrize("deep", [False, True]) +def test_tracking_namespace_copy_preserves_tracking(deep): + import copy + + namespace = TrackingNamespace( + argparse.Namespace(model="example", api_server_count=2), frozenset({"api_server_count"}) + ) + restored = copy.deepcopy(namespace) if deep else copy.copy(namespace) + assert restored.get_explicit_kwargs_dict() == {"api_server_count": 2} + assert restored.model == "example" diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index e7f94ffef..65395265b 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -603,3 +603,21 @@ class TestCumulativeStreamingCoercion: assert all([isinstance(r, SamplingParams) for r in result]) assert result[0].output_kind == RequestOutputKind.DELTA assert result[1].output_kind == RequestOutputKind.DELTA + + +@pytest.mark.parametrize("api_server_count", [1, 2]) +def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count): + from vllm_omni.entrypoints.utils import prepare_stage_config_inputs + + kwargs = { + "api_server_count": api_server_count, + "disable_log_stats": True, + "model_tag": "example", + "max_num_seqs": 8, + "stage_overrides": '{"0": {"max_num_seqs": 4}}', + } + inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None) + + assert inputs.kwargs == {"max_num_seqs": 8} + assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}} + assert kwargs["api_server_count"] == api_server_count diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 7914ecd52..36c85e41d 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -9,6 +9,15 @@ # Profile validated for the c64 / PROMPTS=512 performance experiments: # Stage 0 talker on GPU 0 with S0=64, Stage 1 Code2Wav on GPU 1 with S1=10. # +# For CPU-bound frontend or orchestration workloads, this profile can be used +# with ``--api-server-count 2`` or higher to run multiple API frontends over +# one shared set of local stage engines. This can improve request parsing, +# multimodal/audio preprocessing, and response handling under concurrency; +# it does not duplicate GPU engines or make a single GPU-bound request faster. +# Each frontend keeps process-local state, so runtime voice mutations, sleep / +# wakeup, async video jobs, and resumable duplex sessions are not supported +# with more than one API worker. See docs/cli/serve.md for the full boundary. +# async_chunk: true connectors: diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index d615ebca0..e22bfb241 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -61,11 +61,12 @@ from vllm_omni.engine.stage_client import StageClient from vllm_omni.engine.stage_init_utils import build_stage0_input_processor from vllm_omni.engine.stage_pool import StagePool from vllm_omni.engine.stage_runtime import ( + OmniClientConfig, StageRuntimeInfo, create_stage_runtime, ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import parse_stage_overrides +from vllm_omni.entrypoints.utils import prepare_stage_config_inputs from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter @@ -122,6 +123,7 @@ class OmniEngineBase: _transfer_emitter: Any = None _prom_metrics: Any = None _enable_orch_monitor: bool = False + _client_config: OmniClientConfig | None = None # Lazily created by get_output_blocking_async(). _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None @@ -136,6 +138,7 @@ class OmniEngineBase: log_stats: bool = False, tokenizer: str | None = None, trust_remote_code: bool | None = None, + client_config: OmniClientConfig | None = None, **kwargs: Any, ) -> None: self.model = model @@ -155,6 +158,7 @@ class OmniEngineBase: # --log-stats CLI flag set by the user via OmniBase. self._log_stats = log_stats self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) + self._client_config = client_config logger.info(f"[OmniEngine] Initializing with model {model}") @@ -166,6 +170,8 @@ class OmniEngineBase: _stage_id_kwarg = kwargs.get("stage_id") if isinstance(_stage_id_kwarg, int) and not single_stage_mode: single_stage_mode = True + if client_config is not None and int(client_config.get("client_count", 1)) > 1 and single_stage_mode: + raise ValueError("Multiple API servers cannot be combined with single-stage distributed mode") self.single_stage_mode: bool = single_stage_mode self._single_stage_id_filter: int | None = ( @@ -378,6 +384,7 @@ class OmniEngineBase: omni_lb_policy=self._omni_lb_policy, request_queue=self.request_queue, log_stats=self._log_stats, + client_config=self._client_config, ) self._runtime.initialize() @@ -892,24 +899,15 @@ class OmniEngineBase: ) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" - for legacy_arg in ("stage_configs_path", "stage_configs"): - if legacy_arg in kwargs: - raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") - - # log_stats is captured by __init__; its CLI-only negative alias must - # not cross into per-stage structured config ownership validation. - kwargs.pop("disable_log_stats", None) - deploy_config_path = kwargs.pop("deploy_config", None) - strategy_config_path = kwargs.pop("strategy_config", None) - # CLI callers arrive pre-parsed; offline Python callers may use the - # JSON-string form documented in recipes. - stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) - - # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; - # stage configs know the field as ``streaming_output``. Mirror only a - # truthy value so the CLI's False default does not override deploy YAML. - if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: - kwargs["streaming_output"] = True + config_inputs = prepare_stage_config_inputs( + model, + kwargs, + trust_remote_code=trust_remote_code, + ) + kwargs = config_inputs.kwargs + deploy_config_path = config_inputs.deploy_config_path + strategy_config_path = config_inputs.strategy_config_path + stage_overrides = config_inputs.stage_overrides resolution = cast( _ConfigResolutionResult, diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 612c7481f..1aaebc82b 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -13,6 +13,7 @@ import threading from collections.abc import Callable, Iterator from dataclasses import dataclass, field from multiprocessing import connection +from multiprocessing.process import BaseProcess from types import SimpleNamespace from typing import Any @@ -1138,6 +1139,8 @@ def launch_stage_replica( stage_visible_devices: str | None = None, spawn_device_lock: threading.Lock | None = None, omni_parallel_stage_init: bool = False, + num_api_servers: int = 1, + watched_frontend_processes: list[BaseProcess] | None = None, ) -> Iterator[StageReplicaResources]: """Launch a local LLM stage replica. @@ -1148,6 +1151,8 @@ def launch_stage_replica( keep the same returned resource bundle. """ if omni_master_server is not None: + if num_api_servers != 1: + raise ValueError("Multiple API servers are not supported with distributed stage launch") with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=executor_class, @@ -1173,7 +1178,16 @@ def launch_stage_replica( from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager - addresses = get_engine_zmq_addresses(vllm_config) + if num_api_servers > 1: + # Multi-API Omni is intentionally local-only. IPC paths avoid the + # bind-after-allocation race of tcp://host:0 and are stable for every + # stage, including stages that are not used by APIServerProcessManager. + addresses = EngineZmqAddresses( + inputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)], + outputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)], + ) + else: + addresses = get_engine_zmq_addresses(vllm_config) handshake_address = get_open_zmq_ipc_path() engines_to_handshake = [CoreEngine(index=0, local=True)] with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id): @@ -1197,18 +1211,21 @@ def launch_stage_replica( manager=engine_manager, addresses=addresses, ) + engine_launch = CoreEngineLaunch( + engine_manager=engine_manager, + coordinator=None, + addresses=addresses, + tensor_queue=None, + ) + if watched_frontend_processes is not None: + engine_launch.watched_frontend_processes = watched_frontend_processes wait_for_engine_startup( handshake_socket, engines_to_handshake, vllm_config.parallel_config, False, # coordinated_dp vllm_config.cache_config, - CoreEngineLaunch( - engine_manager=engine_manager, - coordinator=None, - addresses=addresses, - tensor_queue=None, - ), + engine_launch, ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 06f916394..a10cf5861 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1475,6 +1475,8 @@ def build_vllm_config( stage_connector_spec: dict[str, Any] | None = None, engine_args_dict: dict[str, Any] | None = None, headless: bool = False, + api_process_count: int = 1, + api_process_rank: int = 0, ) -> tuple[Any, type]: """Build engine args, then create VllmConfig and executor_class. @@ -1491,6 +1493,9 @@ def build_vllm_config( ) filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict) + if api_process_count != 1 or api_process_rank != 0: + filtered_engine_args_dict["_api_process_count"] = api_process_count + filtered_engine_args_dict["_api_process_rank"] = api_process_rank # _to_dict serializes dataclass fields (e.g. StructuredOutputsConfig) into # plain dicts. When OmniEngineArgs is instantiated with the dict, these @@ -1747,6 +1752,7 @@ def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], stage_init_timeout: int, + locked_devices: set[int] | None = None, ) -> list[int]: """Acquire exclusive file locks on devices needed by this stage. @@ -1801,8 +1807,7 @@ def acquire_device_locks( f"but only {len(physical_devices)} device(s) are available: {physical_devices}" ) - num_devices_to_lock = num_devices_per_stage - devices_to_lock = sorted(physical_devices[:num_devices_to_lock]) + devices_to_lock = sorted(set(physical_devices[:num_devices_per_stage]) - (locked_devices or set())) logger.debug( "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", @@ -1812,7 +1817,7 @@ def acquire_device_locks( prefill_context_parallel_size, sequence_parallel_size, cfg_parallel_size, - num_devices_to_lock, + len(devices_to_lock), devices_to_lock, ) @@ -1830,6 +1835,8 @@ def acquire_device_locks( record_lock_holder_pid(lock_fd, lock_writable) lock_acquired = True lock_fds.append(lock_fd) + if locked_devices is not None: + locked_devices.add(device_id) logger.debug("Acquired exclusive lock for device %s", device_id) except BlockingIOError: os.close(lock_fd) diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index d17e3b08f..ce958ba6d 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -6,14 +6,15 @@ from __future__ import annotations import concurrent.futures +import contextlib import copy import os import threading -import time from collections.abc import Callable, Iterator, Mapping, Sequence -from contextlib import ExitStack, contextmanager -from dataclasses import dataclass -from typing import Any, cast +from contextlib import AbstractContextManager, contextmanager +from dataclasses import dataclass, field +from multiprocessing.process import BaseProcess +from typing import Any, TypedDict, cast import janus from vllm.logger import init_logger @@ -94,6 +95,44 @@ class StageRemoteFactoryContext: executor_class: type | None = None +class OmniClientConfig(TypedDict, total=False): + """Per-frontend transport/configuration passed across process boundaries.""" + + client_count: int + client_index: int + stage_addresses: dict[int | str, dict[int | str, dict[str, Any]]] + input_address: str + output_address: str + stats_update_address: str + actual_address_pipe: Any + tensor_queue: Any + + +@dataclass +class StageEngineLaunch: + """Engine resources handed to a local client or the serving parent. + + Successful launch transfers resource ownership to the caller. Startup or + attachment failures are rolled back by ``launch_stage_engines``. + """ + + client_configs: list[OmniClientConfig] + resources: list[StageReplicaResources] + watched_frontend_processes: list[BaseProcess] = field(default_factory=list) + + def shutdown(self) -> None: + seen: set[int] = set() + for stage_resources in reversed(self.resources): + for resource in (stage_resources.manager, stage_resources.coordinator): + if resource is None or id(resource) in seen: + continue + seen.add(id(resource)) + try: + resource.shutdown() + except Exception: + logger.warning("[StageEngineLaunch] resource shutdown failed", exc_info=True) + + def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: try: normalized = LoadBalancingPolicy(policy) @@ -125,13 +164,14 @@ class StageRuntime: self, stage_configs: list[Any], model: str, - config_path: str, + config_path: str | None, *, stage_init_timeout: int, async_chunk: bool, tokenizer: str | None = None, parallel_stage_init: bool = False, log_stats: bool = False, + client_config: OmniClientConfig | None = None, ) -> None: self._stage_configs = stage_configs self._model = model @@ -146,6 +186,14 @@ class StageRuntime: self._parallel_stage_init = parallel_stage_init self._log_stats = log_stats self._num_stages = len(stage_configs) + self._client_count = int((client_config or {}).get("client_count", 1)) + self._client_index = int((client_config or {}).get("client_index", 0)) + if self._client_count < 1: + raise ValueError(f"client_count must be >= 1, got {self._client_count}") + if self._client_index != -1 and not 0 <= self._client_index < self._client_count: + raise ValueError(f"client_index must be in [0, {self._client_count}), got {self._client_index}") + self._external_stage_addresses = (client_config or {}).get("stage_addresses") + self._api_process_rank = self._client_index # Populated by initialize() self.stage_pools: list[StagePool] = [] @@ -264,6 +312,205 @@ class StageRuntime: self._cleanup_after_initialize_failure() raise exc + @contextmanager + def launch_stage_engines( + self, + num_api_servers: int = 1, + *, + stage_plans: Sequence[LogicalStageInitPlan] | None = None, + stage_init_timeout: int | None = None, + ) -> Iterator[StageEngineLaunch]: + """Launch EngineCore stages for local or subprocess client attachment. + + The regular runtime supplies a replica plan from its device-group + scheduler; the serving parent resolves the whole pipeline here. Both + use the same device locks, spawn, readiness and rollback machinery. + One client completes readiness before attachment, preserving the + existing in-process frontend. Multiple clients attach in the caller's + ``with`` body before readiness completes on context exit. + """ + if self._external_stage_addresses is not None: + raise ValueError("A runtime attached to external stage engines cannot launch engines") + if num_api_servers < 1: + raise ValueError(f"num_api_servers must be >= 1, got {num_api_servers}") + timeout = self._stage_init_timeout if stage_init_timeout is None else stage_init_timeout + resolve_plans = stage_plans is None + if stage_plans is None: + self._client_count = num_api_servers + self._api_process_rank = -1 if num_api_servers > 1 else self._client_index + self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var) + stage_plans = self._prepare_stage_plans() + + unsupported = [ + plan.stage_id + for plan in stage_plans + if any(replica.metadata.stage_type == "diffusion" for replica in plan.replicas) + ] + if unsupported: + raise ValueError( + "--api-server-count currently supports EngineCore stages only; " + f"diffusion stage(s) are not supported: {unsupported}" + ) + if any(replica.launch_mode != "local" for plan in stage_plans for replica in plan.replicas): + raise ValueError("--api-server-count is not supported with remote or headless stages") + + if num_api_servers > 1: + parallel_configs: list[Any] = [] + for plan in stage_plans: + for replica in plan.replicas: + if replica.stage_vllm_config is None or replica.executor_class is None: + raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration") + parallel_configs.append(replica.stage_vllm_config.parallel_config) + if any(getattr(parallel_config, "enable_fault_tolerance", False) for parallel_config in parallel_configs): + raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance") + if any(getattr(parallel_config, "enable_elastic_ep", False) for parallel_config in parallel_configs): + raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep") + if any( + getattr(parallel_config, "data_parallel_size", 1) != 1 or getattr(parallel_config, "use_ray", False) + for parallel_config in parallel_configs + ): + raise ValueError( + "--api-server-count > 1 currently supports one local process group per replica; " + "intra-replica data parallelism and Ray backends are not supported" + ) + + if resolve_plans and self._parallel_stage_init: + self._reject_unguardable_executors(stage_plans) + self._run_stage_admission(stage_plans) + + client_configs: list[OmniClientConfig] = [ + OmniClientConfig( + client_count=num_api_servers, + client_index=client_index, + stage_addresses={}, + ) + for client_index in range(num_api_servers) + ] + entered_contexts: list[AbstractContextManager[StageReplicaResources]] = [] + resources: list[StageReplicaResources] = [] + watched_frontend_processes: list[BaseProcess] = [] + # launch_stage_replica performs its readiness handshake in __exit__, + # after API workers are started, so initialization locks must remain + # held until all launch contexts have exited. + lock_fds: list[int] = [] + locked_devices: set[int] = set() + launch: StageEngineLaunch | None = None + exited_contexts = 0 + + try: + for plan in stage_plans: + for replica in plan.replicas: + if replica.stage_vllm_config is None or replica.executor_class is None: + raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration") + if replica.engine_args_dict is None: + raise RuntimeError(f"LLM stage {plan.stage_id} is missing engine args") + + physical_devices = self._resolve_replica_physical_devices( + replica.metadata.stage_id, + replica.metadata.runtime_cfg, + ) + if not self._parallel_stage_init: + with self._scoped_spawn_device_env(physical_devices): + lock_fds.extend( + acquire_device_locks( + replica.metadata.stage_id, + replica.engine_args_dict, + timeout, + locked_devices, + ) + ) + + launch_context = launch_stage_replica( + vllm_config=replica.stage_vllm_config, + executor_class=replica.executor_class, + log_stats=self._log_stats, + stage_id=replica.metadata.stage_id, + replica_id=replica.replica_id, + stage_config=replica.stage_cfg, + omni_master_server=self._get_omni_master_server(), + omni_coordinator_address=self._get_coordinator_address(), + stage_visible_devices=physical_devices, + spawn_device_lock=self._spawn_device_lock, + omni_parallel_stage_init=self._parallel_stage_init, + num_api_servers=num_api_servers, + watched_frontend_processes=watched_frontend_processes if num_api_servers > 1 else None, + ) + # Environment overlays are process-global. Serialize spawn; + # parallel initialization waits for READY outside this lock. + with self._replica_launch_lock: + with stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg): + stage_resources = launch_context.__enter__() + entered_contexts.append(launch_context) + if stage_resources is None: + raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no resources") + resources.append(stage_resources) + if num_api_servers == 1: + exited_contexts += 1 + launch_context.__exit__(None, None, None) + addresses = stage_resources.addresses + # The generic resource bundle allows no addresses (e.g. + # diffusion), but EngineCore attachment requires them. + if addresses is None: + raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no addresses") + # Unlike the single-client handshake, only stage 0 is + # handed to APIServerProcessManager for address updates. + # Other stages must already expose connectable endpoints. + if num_api_servers > 1 and any( + address.startswith("tcp://") and address.rsplit(":", 1)[-1] == "0" + for address in (*addresses.inputs, *addresses.outputs) + ): + raise RuntimeError( + f"Stage {plan.stage_id} returned deferred TCP addresses; " + "multi-API launch requires fixed ports or IPC addresses" + ) + if len(addresses.inputs) != num_api_servers or len(addresses.outputs) != num_api_servers: + raise RuntimeError( + f"Stage {plan.stage_id} returned {len(addresses.inputs)} input and " + f"{len(addresses.outputs)} output addresses for {num_api_servers} API servers" + ) + + if num_api_servers == 1 and lock_fds: + release_device_locks(lock_fds) + lock_fds.clear() + locked_devices.clear() + + for client_index, client_config in enumerate(client_configs): + stage_addresses = client_config["stage_addresses"] + replica_addresses = stage_addresses.setdefault(plan.stage_id, {}) + replica_addresses[replica.replica_id] = { + "input_address": addresses.inputs[client_index], + "output_address": addresses.outputs[client_index], + } + if addresses.frontend_stats_publish_address is not None: + replica_addresses[replica.replica_id]["stats_update_address"] = ( + addresses.frontend_stats_publish_address + ) + + launch = StageEngineLaunch( + client_configs=client_configs, + resources=resources, + watched_frontend_processes=watched_frontend_processes, + ) + yield launch + + while exited_contexts < len(entered_contexts): + entered_context = entered_contexts[exited_contexts] + exited_contexts += 1 + entered_context.__exit__(None, None, None) + except BaseException as exc: + if launch is not None: + launch.shutdown() + else: + StageEngineLaunch(client_configs=client_configs, resources=resources).shutdown() + exc_info = (type(exc), exc, exc.__traceback__) + for pending_context in reversed(entered_contexts[exited_contexts:]): + with contextlib.suppress(BaseException): + pending_context.__exit__(*exc_info) + raise + finally: + if lock_fds: + release_device_locks(lock_fds) + def shutdown(self) -> None: for pool in self.stage_pools: for client in pool.clients: @@ -406,6 +653,8 @@ class StageRuntime: self._model, stage_connector_spec=stage_connector_spec, engine_args_dict=engine_args_dict, + api_process_count=self._client_count, + api_process_rank=self._api_process_rank, ) for replica_id in range(num_replicas): @@ -737,122 +986,70 @@ class StageRuntime: stage_init_timeout: int, ) -> StagePoolClient: """Initialize one local LLM replica using vLLM's launch/attach pattern.""" - resources: StageReplicaResources | None = None - stage_client = None - lock_fds: list[int] = [] - try: - physical_devices = self._resolve_replica_physical_devices( - plan.metadata.stage_id, - plan.metadata.runtime_cfg, - ) - if physical_devices: - logger.info( - "[stage_init] Stage-%s set runtime devices: %s", - plan.metadata.stage_id, - physical_devices, - ) - vllm_config = plan.stage_vllm_config - executor_class = plan.executor_class - if vllm_config is None: - raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing vllm_config") - if executor_class is None: - raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class") - if plan.engine_args_dict is None: - raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args") - # G3 device locks: in the default (serial) path the parent holds a - # per-device LOCK_EX across the whole child init. When parallel stage - # init is enabled the engine-core child takes phased SH/EX locks - # itself (see stage_phase_lock), so the parent must NOT also hold the - # lock here — it would self-deadlock while waiting for the child's - # READY handshake. - if not self._parallel_stage_init: - g3_start = time.perf_counter() - with self._scoped_spawn_device_env(physical_devices): - lock_fds = acquire_device_locks( - plan.metadata.stage_id, - plan.engine_args_dict, - stage_init_timeout, - ) - logger.debug( - "[stage_init] Stage-%s G3 device-lock acquire took %.3fs", - plan.metadata.stage_id, - time.perf_counter() - g3_start, - ) - - launch_cm = launch_stage_replica( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - stage_id=plan.metadata.stage_id, - replica_id=plan.replica_id, - stage_config=plan.stage_cfg, - omni_master_server=self._get_omni_master_server(), - omni_coordinator_address=self._get_coordinator_address(), - stage_visible_devices=physical_devices, - spawn_device_lock=self._spawn_device_lock, - omni_parallel_stage_init=self._parallel_stage_init, - ) - # Only process spawning and the runtime.env overlay need the global - # launch lock. The launch context's exit waits for READY: holding the - # lock there serializes model loading and compilation even across - # different GPUs. Default initialization keeps its per-device EX - # locks until READY; parallel_stage_init uses child phase locks. - g2_start = time.perf_counter() - with ExitStack() as launch_stack: - with self._replica_launch_lock: - g2_locked = time.perf_counter() - with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - resources = launch_stack.enter_context(launch_cm) - g2_spawned = time.perf_counter() - logger.debug( - "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs", - plan.metadata.stage_id, - g2_locked - g2_start, - g2_spawned - g2_locked, - time.perf_counter() - g2_spawned, + external_addresses = self._get_external_client_addresses( + plan.metadata.stage_id, + plan.replica_id, + ) + if external_addresses is not None: + if plan.stage_vllm_config is None or plan.executor_class is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration") + return cast( + StagePoolClient, + StageEngineCoreClientBase.make_async_mp_client( + vllm_config=plan.stage_vllm_config, + executor_class=plan.executor_class, + log_stats=self._log_stats, + metadata=plan.metadata, + client_addresses=external_addresses, + client_count=self._client_count, + client_index=self._client_index, + ), ) - logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id) - if resources is None: - raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no resources") - if resources.addresses is None: - raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no addresses") - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - metadata=plan.metadata, - client_addresses=self._client_addresses_from_zmq(resources.addresses), - engine_manager=resources.manager, - coordinator=resources.coordinator, + if plan.stage_vllm_config is None or plan.executor_class is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration") + stage_plan = LogicalStageInitPlan( + stage_idx=plan.metadata.stage_id, + stage_id=plan.metadata.stage_id, + replicas=[plan], + ) + with self.launch_stage_engines( + stage_plans=[stage_plan], + stage_init_timeout=stage_init_timeout, + ) as launch: + resources = launch.resources[0] + return cast( + StagePoolClient, + StageEngineCoreClientBase.make_async_mp_client( + vllm_config=plan.stage_vllm_config, + executor_class=plan.executor_class, + log_stats=self._log_stats, + metadata=plan.metadata, + client_addresses=self._client_addresses_from_zmq(resources.addresses), + engine_manager=resources.manager, + coordinator=resources.coordinator, + ), ) - logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id) - return cast(StagePoolClient, stage_client) - except Exception: - if stage_client is not None: - try: - stage_client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[StageRuntime] Failed to cleanup stage %s after init failure: %s", - plan.metadata.stage_id, - cleanup_error, - ) - else: - self._cleanup_launched_resources( - stage_id=plan.metadata.stage_id, - resources=resources, - ) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - def _get_coordinator_address(self) -> str | None: """Return coordinator router address. Overridden by DistStageRuntime.""" return None + def _get_external_client_addresses(self, stage_id: int, replica_id: int) -> dict[str, Any] | None: + if self._external_stage_addresses is None: + return None + stage_addresses = self._external_stage_addresses.get(stage_id) + if stage_addresses is None: + stage_addresses = self._external_stage_addresses.get(str(stage_id)) + if stage_addresses is None: + raise ValueError(f"Missing external addresses for stage {stage_id}") + replica_addresses = stage_addresses.get(replica_id) + if replica_addresses is None: + replica_addresses = stage_addresses.get(str(replica_id)) + if replica_addresses is None: + raise ValueError(f"Missing external addresses for stage {stage_id} replica {replica_id}") + return dict(replica_addresses) + def _get_omni_master_server(self) -> OmniMasterServer | None: """Return the master server for local distributed launches, if any.""" return None @@ -991,7 +1188,7 @@ class DistStageRuntime(StageRuntime): self, stage_configs: list[Any], model: str, - config_path: str, + config_path: str | None, *, stage_init_timeout: int, async_chunk: bool, @@ -1324,7 +1521,7 @@ class DistStageRuntime(StageRuntime): def create_stage_runtime( stage_configs: list[Any], model: str, - config_path: str, + config_path: str | None, *, single_stage_mode: bool, stage_init_timeout: int, @@ -1340,6 +1537,7 @@ def create_stage_runtime( omni_lb_policy: str = "random", request_queue: janus.Queue[EngineQueueMessage] | None = None, log_stats: bool = False, + client_config: OmniClientConfig | None = None, ) -> StageRuntime: """Factory: select StageRuntime or DistStageRuntime.""" if single_stage_mode: @@ -1371,4 +1569,5 @@ def create_stage_runtime( tokenizer=tokenizer, parallel_stage_init=parallel_stage_init, log_stats=log_stats, + client_config=client_config, ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index a35f06202..493f2533e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -8,13 +8,16 @@ Supports both multi-stage LLM models (e.g., Qwen2.5-Omni) and diffusion models (e.g., Qwen-Image) through the same CLI interface. """ +from __future__ import annotations + import argparse +import contextlib import json import math import os import signal from types import FrameType -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import uvloop from vllm.entrypoints.cli.types import CLISubcommand @@ -23,10 +26,18 @@ from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger from vllm_omni.entrypoints.cli.logo import log_logo -from vllm_omni.entrypoints.openai.api_server import omni_run_server -from vllm_omni.entrypoints.utils import parse_stage_overrides +from vllm_omni.entrypoints.openai.api_server import ( + omni_run_server, + run_omni_api_server_worker_proc, +) +from vllm_omni.entrypoints.utils import parse_stage_overrides, prepare_stage_config_inputs from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace +if TYPE_CHECKING: + from vllm.v1.utils import APIServerProcessManager + + from vllm_omni.engine.stage_runtime import StageEngineLaunch, StageRuntime + logger = init_logger(__name__) DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models @@ -141,6 +152,8 @@ class OmniServeCommand(CLISubcommand): if args.headless: run_headless(args) + elif (getattr(args, "api_server_count", None) or 1) > 1: + run_multi_api_server_omni(args) else: uvloop.run(omni_run_server(args)) @@ -198,6 +211,39 @@ class OmniServeCommand(CLISubcommand): if getattr(args, "omni_replica_address", None) is not None and not args.headless: raise ValueError("--omni-replica-address requires --headless to be set") + api_server_count = getattr(args, "api_server_count", None) + if api_server_count is not None and api_server_count < 1 and not args.headless: + raise ValueError("--api-server-count must be >= 1 unless --headless is set") + if args.headless and api_server_count is not None and api_server_count > 0: + raise ValueError("--api-server-count cannot be used with --headless") + if api_server_count is not None and api_server_count > 1: + distributed_args = ( + "stage_id", + "omni_master_address", + "omni_master_port", + "omni_replica_address", + ) + if any(getattr(args, name, None) is not None for name in distributed_args): + raise ValueError( + "--api-server-count > 1 cannot be combined with stage-based " + "or distributed stage arguments (--stage-id/--omni-master-* / " + "--omni-replica-address)" + ) + if getattr(args, "omni_dp_size_local", 1) != 1: + raise ValueError("--api-server-count > 1 requires --omni-dp-size-local=1") + if getattr(args, "worker_backend", "multi_process") != "multi_process": + raise ValueError("--api-server-count > 1 requires --worker-backend=multi_process") + if getattr(args, "enable_fault_tolerance", False): + raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance") + if getattr(args, "enable_elastic_ep", False): + raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep") + if getattr(args, "enable_sleep_mode", False): + raise ValueError("--api-server-count > 1 cannot be combined with --enable-sleep-mode") + from vllm import envs as vllm_envs + + if getattr(vllm_envs, "VLLM_ALLOW_RUNTIME_LORA_UPDATING", False): + raise ValueError("--api-server-count > 1 cannot be combined with VLLM_ALLOW_RUNTIME_LORA_UPDATING") + # --omni-dp-size-local is process-local. A value other than 1 only # makes sense when this process owns a stage (head or headless). omni_dp_size_local = getattr(args, "omni_dp_size_local", None) @@ -221,7 +267,6 @@ class OmniServeCommand(CLISubcommand): "data_parallel_rpc_port": "--data-parallel-rpc-port", "data_parallel_start_rank": "--data-parallel-start-rank", "data_parallel_backend": "--data-parallel-backend", - "api_server_count": "--api-server-count", "enable_expert_parallel": "--enable-expert-parallel", } offenders = sorted(flag for dest, flag in prohibited_with_omni.items() if dest in explicit_cli_keys) @@ -261,6 +306,8 @@ class OmniServeCommand(CLISubcommand): model = getattr(args, "model_tag", None) or getattr(args, "model", None) if model and is_diffusion_model(model): + if api_server_count is not None and api_server_count > 1: + raise ValueError("--api-server-count > 1 is not supported for diffusion models") logger.info("Detected diffusion model: %s", model) return validate_parsed_serve_args(args) @@ -1027,6 +1074,207 @@ class OmniServeCommand(CLISubcommand): return serve_parser +def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int) -> StageRuntime: + """Resolve the local EngineCore stages that the parent process owns.""" + from vllm_omni.config.resolver import resolve_omni_config + from vllm_omni.engine.stage_runtime import StageRuntime + + kwargs = args.get_explicit_kwargs_dict() + model = kwargs.pop("model", None) or args.model + + trust_remote_code = kwargs.get("trust_remote_code") + if trust_remote_code is False: + trust_remote_code = None + config_inputs = prepare_stage_config_inputs( + model, + kwargs, + trust_remote_code=trust_remote_code, + snapshot_model=True, + ) + model = config_inputs.model + kwargs = config_inputs.kwargs + resolved = resolve_omni_config( + model, + cli_overrides=kwargs, + trust_remote_code=config_inputs.trust_remote_code, + deploy_config_path=config_inputs.deploy_config_path, + stage_overrides=config_inputs.stage_overrides, + strategy_config_path=config_inputs.strategy_config_path, + ) + + config_path = resolved.config_path + stage_configs = list(resolved.stage_configs) + + sleep_stages = [ + int(getattr(stage_config, "stage_id", stage_index)) + for stage_index, stage_config in enumerate(stage_configs) + if bool( + getattr( + getattr(stage_config, "model_config", getattr(stage_config, "engine_args", None)), + "enable_sleep_mode", + False, + ) + ) + ] + if sleep_stages: + raise ValueError( + "--api-server-count > 1 cannot be combined with sleep mode; " + f"disable enable_sleep_mode for stage(s) {sleep_stages}" + ) + + async_chunk = any( + bool( + getattr( + getattr(stage, "connector_config", None), + "async_chunk", + getattr(getattr(stage, "engine_args", None), "async_chunk", False), + ) + ) + for stage in stage_configs + ) + return StageRuntime( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=int(getattr(args, "stage_init_timeout", 300)), + parallel_stage_init=bool(getattr(args, "parallel_stage_init", False)), + async_chunk=async_chunk, + tokenizer=getattr(args, "tokenizer", None), + log_stats=not bool(getattr(args, "disable_log_stats", False)), + ) + + +def _wait_for_multi_api_server_completion( + api_server_manager: APIServerProcessManager, + engine_launch: StageEngineLaunch, +) -> None: + """Wait until API workers complete or any shared stage engine fails.""" + from multiprocessing import connection + + api_processes = list(api_server_manager.processes) + engine_processes = [ + process + for resources in engine_launch.resources + if resources.manager is not None + for process in resources.manager.processes + ] + api_by_sentinel = {process.sentinel: process for process in api_processes} + engine_by_sentinel = {process.sentinel: process for process in engine_processes} + + while api_by_sentinel: + ready = connection.wait([*api_by_sentinel, *engine_by_sentinel]) + for sentinel in ready: + if sentinel in engine_by_sentinel: + process = engine_by_sentinel[sentinel] + raise RuntimeError( + f"Shared stage engine process {process.name} (PID: {process.pid}) " + f"exited with code {process.exitcode}" + ) + process = api_by_sentinel.pop(sentinel) + if process.exitcode != 0: + raise RuntimeError( + f"API server process {process.name} (PID: {process.pid}) exited with code {process.exitcode}" + ) + + +def _start_api_server_process_manager( + *, + cleanup_timeout: float, + **manager_kwargs: object, +) -> APIServerProcessManager: + """Construct vLLM's manager with rollback for a partial ``__init__``. + + ``APIServerProcessManager`` starts workers one by one in ``__init__`` and + installs its finalizer only after every ``Process.start`` succeeds. Keep a + reference to the partially initialized object so workers started before a + later failure are still terminated. + """ + from vllm.v1.utils import APIServerProcessManager, shutdown + + manager = APIServerProcessManager.__new__(APIServerProcessManager) + try: + APIServerProcessManager.__init__(manager, **manager_kwargs) + except BaseException: + try: + for pipe in getattr(manager, "_address_pipes", ()): + with contextlib.suppress(Exception): + pipe.close() + # A Process whose start() raised is still present in the upstream + # manager's list, but calling is_alive() on it raises because it + # has no Popen object. Only hand successfully started children to + # vLLM's shutdown helper. + started_processes = [process for process in getattr(manager, "processes", ()) if process.pid is not None] + if started_processes: + shutdown(started_processes, timeout=cleanup_timeout) + except Exception: + logger.exception("Failed to clean up partially started API server processes") + raise + return manager + + +def run_multi_api_server_omni(args: TrackingNamespace) -> None: + """Launch API subprocesses that share one set of local stage engines.""" + from vllm.entrypoints.openai.api_server import setup_server + from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus + + num_api_servers = int(args.api_server_count) + if num_api_servers < 2: + raise ValueError(f"api_server_count must be >= 2, got {num_api_servers}") + + setup_multiprocess_prometheus() + shutdown_requested = False + + def signal_handler(signum: int, frame: FrameType | None) -> None: + nonlocal shutdown_requested + logger.debug("Received %d signal.", signum) + if not shutdown_requested: + shutdown_requested = True + raise SystemExit + + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + + listen_address, sock = setup_server(args, reuse_port=True) + stage_runtime = None + api_server_manager = None + engine_launch = None + try: + stage_runtime = _build_multi_api_stage_runtime(args, num_api_servers) + with stage_runtime.launch_stage_engines(num_api_servers) as engine_launch: + args._omni_stage_client_configs = engine_launch.client_configs + primary_addresses = engine_launch.resources[0].addresses + if primary_addresses is None: + raise RuntimeError("Primary stage engine returned no addresses") + timeout = float(getattr(args, "shutdown_timeout", 5) or 5) + api_server_manager = _start_api_server_process_manager( + cleanup_timeout=timeout, + listen_address=listen_address, + sock=sock, + args=args, + num_servers=num_api_servers, + input_addresses=primary_addresses.inputs, + output_addresses=primary_addresses.outputs, + target_server_fn=run_omni_api_server_worker_proc, + ) + engine_launch.watched_frontend_processes.extend(api_server_manager.processes) + actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses() + primary_addresses.inputs = actual_inputs + primary_addresses.outputs = actual_outputs + + _wait_for_multi_api_server_completion(api_server_manager, engine_launch) + finally: + timeout = float(getattr(args, "shutdown_timeout", 5) or 5) + if api_server_manager is not None: + api_server_manager.shutdown(timeout=timeout) + if engine_launch is not None: + engine_launch.shutdown() + if stage_runtime is not None: + stage_runtime.shutdown() + sock.close() + if shutdown_requested: + logger.info("Shared API server shutdown completed") + + def run_headless(args: TrackingNamespace) -> None: """Run a single stage in headless mode. @@ -1075,15 +1323,6 @@ def run_headless(args: TrackingNamespace) -> None: # Filter down to a dict of things explicitly requested by the user args_dict = args.get_explicit_kwargs_dict() - # This controls only the API-process WebSocket and is not a stage-engine option. - args_dict.pop("robot_openpi_idle_timeout", None) - # This CLI-only negative alias is consumed below when selecting the - # launcher log_stats value; it is not a per-stage config override. - args_dict.pop("disable_log_stats", None) - - deploy_config_path = args_dict.pop("deploy_config", None) - strategy_config_path = args_dict.pop("strategy_config", None) - stage_overrides = args_dict.pop("stage_overrides", None) # ``--replica-id`` is deprecated and ignored — replica ids are # auto-assigned by ``OmniMasterServer`` so headless processes carry @@ -1099,15 +1338,21 @@ def run_headless(args: TrackingNamespace) -> None: ) args_dict.pop("replica_id") - resolved = resolve_omni_config( + config_inputs = prepare_stage_config_inputs( model, + args_dict, # store_true cannot express an explicit False: absent maps to None # ("not specified") so the deploy yaml's per-stage value applies. trust_remote_code=getattr(args, "trust_remote_code", None) or None, + ) + args_dict = config_inputs.kwargs + resolved = resolve_omni_config( + model, + trust_remote_code=config_inputs.trust_remote_code, cli_overrides=args_dict, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, + deploy_config_path=config_inputs.deploy_config_path, + stage_overrides=config_inputs.stage_overrides, + strategy_config_path=config_inputs.strategy_config_path, ) config_path = resolved.config_path stage_configs = list(resolved.stage_configs) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 20706396c..94f1ab513 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -6,6 +6,7 @@ This module owns app construction, server startup, app-state initialization, and route bodies that have not yet moved to endpoint-owned modules.""" import asyncio +import copy import dataclasses import json import multiprocessing @@ -14,6 +15,8 @@ import os # Image generation API imports import random +import signal +import socket import time from argparse import Namespace from collections.abc import AsyncIterator @@ -22,6 +25,7 @@ from http import HTTPStatus from pathlib import Path from typing import Annotated, Any, Literal +import uvloop import vllm.envs as envs from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse @@ -81,12 +85,14 @@ from vllm.renderers.online_renderer import OnlineRenderer from vllm.tasks import POOLING_TASKS from vllm.tool_parsers import ToolParserManager from vllm.utils import random_uuid -from vllm.utils.system_utils import decorate_logs +from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) +from vllm_omni.engine.stage_init_utils import set_death_signal +from vllm_omni.engine.stage_runtime import OmniClientConfig from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime @@ -232,8 +238,55 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) -async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None: +def run_omni_api_server_worker_proc( + listen_address: str, + sock: socket.socket, + args: TrackingNamespace, + client_config: dict[str, Any] | None = None, + **uvicorn_kwargs: object, +) -> None: + """Entrypoint used by vLLM's API server process manager.""" + set_death_signal(signal.SIGTERM) + manager_config = client_config or {} + client_index = int(manager_config.get("client_index", 0)) + all_client_configs = getattr(args, "_omni_stage_client_configs", None) + if not all_client_configs or not 0 <= client_index < len(all_client_configs): + raise RuntimeError(f"Missing Omni stage client configuration for API server {client_index}") + + omni_client_config: OmniClientConfig = copy.deepcopy(all_client_configs[client_index]) + omni_client_config["client_count"] = int(manager_config.get("client_count", 1)) + omni_client_config["client_index"] = client_index + + stage_addresses = omni_client_config["stage_addresses"] + first_stage_id = min(stage_addresses) + first_replica_id = min(stage_addresses[first_stage_id]) + first_addresses = stage_addresses[first_stage_id][first_replica_id] + for key in ("input_address", "output_address", "actual_address_pipe", "tensor_queue"): + if key in manager_config: + first_addresses[key] = manager_config[key] + + set_process_title("APIServer", str(client_index)) + decorate_logs("APIServer", skip_if_decorated=True) + uvloop.run( + omni_run_server_worker( + listen_address, + sock, + args, + client_config=omni_client_config, + **uvicorn_kwargs, + ) + ) + + +async def omni_run_server_worker( + listen_address: str, + sock: socket.socket, + args: TrackingNamespace, + client_config: OmniClientConfig | None = None, + **uvicorn_kwargs: object, +) -> None: """Run a single API server worker.""" + api_server_count = _resolve_api_server_count(args, client_config) if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: ToolParserManager.import_tool_parser(args.tool_parser_plugin) @@ -273,6 +326,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, # OMNI: Pass supported_tasks to build_app (required by upstream vLLM) app = build_openai_app(args, supported_tasks) + app.state.api_server_count = api_server_count # OMNI: Remove upstream routes that we override with omni-specific handlers remove_route_from_app(app, "/v1/chat/completions", {"POST"}) @@ -348,6 +402,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, if scope["type"] == "http": scope.setdefault("state", {}) scope["state"]["request_timestamp"] = time.time() + await self._inner(scope, receive, send) # Startup duplex warmup (duplex_session.warmup_frames in the deploy @@ -406,7 +461,7 @@ async def build_async_omni( args: TrackingNamespace, *, disable_frontend_multiprocessing: bool | None = None, - client_config: dict[str, Any] | None = None, + client_config: OmniClientConfig | None = None, ) -> AsyncIterator[EngineClient]: """Build an AsyncOmni instance from command-line arguments. @@ -437,6 +492,7 @@ async def build_async_omni( async with build_async_omni_from_stage_config( args, disable_frontend_multiprocessing=disable_frontend_multiprocessing, + client_config=client_config, ) as async_omni: yield async_omni @@ -476,6 +532,7 @@ async def build_async_omni_from_stage_config( args: TrackingNamespace, *, disable_frontend_multiprocessing: bool = False, + client_config: OmniClientConfig | None = None, ) -> AsyncIterator[EngineClient]: """Create AsyncOmni from stage configuration. @@ -532,6 +589,8 @@ async def build_async_omni_from_stage_config( kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) + if client_config is not None: + kwargs["client_config"] = client_config if _should_serve_duplex(model, kwargs): async_omni = DuplexOmni(model=model, **kwargs) else: @@ -1430,6 +1489,42 @@ async def list_voices(raw_request: Request): return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers}) +def _resolve_api_server_count(args: Namespace, client_config: OmniClientConfig | None) -> int: + """Resolve frontend topology once, from the launcher's client configuration.""" + requested = getattr(args, "api_server_count", None) + actual = 1 if client_config is None else client_config.get("client_count") + if type(actual) is not int or actual < 1: + raise ValueError("API client configuration requires a positive integer client_count") + if requested is not None and (type(requested) is not int or requested < 1 or requested != actual): + raise ValueError("API server count does not match the frontend launch configuration") + return actual + + +def _reject_process_local_state_with_multiple_api_workers(raw_request: Request, operation: str) -> None: + """Reject access to frontend state that is not shared across API workers.""" + api_server_count = getattr(raw_request.app.state, "api_server_count", None) + if type(api_server_count) is not int or api_server_count < 1: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="API worker topology is not initialized.", + ) + if api_server_count > 1: + raise HTTPException( + status_code=HTTPStatus.CONFLICT.value, + detail=( + f"{operation} is not supported with --api-server-count > 1 because " + "the operation uses process-local frontend state" + ), + ) + + +async def _require_single_api_video_store(raw_request: Request) -> None: + # VIDEO_STORE and VIDEO_TASKS are process-local, including reads and + # cancellation. Keep this independent of the temporary diffusion launch + # restriction: enabling multi-API diffusion alone cannot make jobs shared. + _reject_process_local_state_with_multiple_api_workers(raw_request, "Asynchronous video jobs") + + @router.post( "/v1/audio/voices", responses={ @@ -1475,6 +1570,7 @@ async def upload_voice( Returns: JSON response with voice information """ + _reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice upload") handler = Omnispeech(raw_request) if handler is None: return _create_speech_error_json_response( @@ -1540,6 +1636,7 @@ async def delete_voice(name: str, raw_request: Request): Returns: JSON response indicating success or failure """ + _reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice deletion") handler = Omnispeech(raw_request) if handler is None: return _create_speech_error_json_response( @@ -1633,13 +1730,15 @@ async def streaming_video_output(websocket: WebSocket): @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" - await _wait_for_duplex_warmup(websocket) duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None) duplex_query = websocket.query_params.get("duplex") use_duplex_realtime = duplex_handler is not None and ( duplex_query is None or (isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"}) ) if use_duplex_realtime and duplex_handler is not None: + if await _reject_multi_api_duplex(websocket): + return + await _wait_for_duplex_warmup(websocket) await duplex_handler.handle_realtime_session(websocket) return @@ -1669,6 +1768,8 @@ async def _wait_for_duplex_warmup(websocket: WebSocket) -> None: @router.websocket("/v1/duplex") async def duplex_websocket(websocket: WebSocket): """Alias of ``/v1/realtime?duplex=1``: the same Realtime duplex session protocol.""" + if await _reject_multi_api_duplex(websocket): + return await _wait_for_duplex_warmup(websocket) handler = getattr(websocket.app.state, "openai_serving_duplex", None) if handler is None: @@ -1709,6 +1810,24 @@ async def realtime_robot_openpi(websocket: WebSocket): await connection.handle_connection() +async def _reject_multi_api_duplex(websocket: WebSocket) -> bool: + # Session credentials and replay state are local to each frontend; a shared + # listening socket cannot route reconnects back to the session owner. + count = getattr(websocket.app.state, "api_server_count", None) + if count is None or (type(count) is int and count == 1): + return False + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "error": "Duplex sessions require a single API worker with initialized topology.", + "code": "multi_api_duplex_unsupported", + } + ) + await websocket.close(code=1008) + return True + + # RL Rollout serving (RFC #3747, P0) @@ -2476,6 +2595,7 @@ async def edit_images( @router.post( "/v1/videos", + dependencies=[Depends(_require_single_api_video_store)], responses={ HTTPStatus.OK.value: {"model": VideoResponse}, HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, @@ -2617,7 +2737,7 @@ async def create_video_sync( ) -@router.get("/v1/videos", response_model=VideoListResponse) +@router.get("/v1/videos", response_model=VideoListResponse, dependencies=[Depends(_require_single_api_video_store)]) async def list_videos( after: str | None = None, limit: int | None = Query(None, ge=0, le=100), @@ -2654,7 +2774,7 @@ async def list_videos( return VideoListResponse(data=jobs, has_more=has_more, first_id=first_id, last_id=last_id) -@router.get("/v1/videos/{video_id}", response_model=None) +@router.get("/v1/videos/{video_id}", response_model=None, dependencies=[Depends(_require_single_api_video_store)]) async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse: """Retrieve metadata for a previously created video job. @@ -2682,7 +2802,7 @@ async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse: return job -@router.delete("/v1/videos/{video_id}") +@router.delete("/v1/videos/{video_id}", dependencies=[Depends(_require_single_api_video_store)]) async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteResponse: """Delete a stored video job and any generated output. @@ -2752,7 +2872,7 @@ async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteRespon return VideoDeleteResponse(id=job.id, deleted=True) -@router.get("/v1/videos/{video_id}/content") +@router.get("/v1/videos/{video_id}/content", dependencies=[Depends(_require_single_api_video_store)]) async def download_video(video_id: str) -> Response: """Download the generated file for a completed video job. @@ -2845,6 +2965,7 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No @router.post("/v1/omni/sleep") async def omni_sleep(request: OmniSleepRequest, raw_request: Request): + _reject_process_local_state_with_multiple_api_workers(raw_request, "Sleep") engine_client = raw_request.app.state.engine_client sleeping_set = raw_request.app.state.sleeping_stages if not hasattr(engine_client, "sleep"): @@ -2857,6 +2978,7 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request): @router.post("/v1/omni/wakeup") async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): + _reject_process_local_state_with_multiple_api_workers(raw_request, "Wakeup") engine_client = raw_request.app.state.engine_client sleeping_set = raw_request.app.state.sleeping_stages if not any(sid in sleeping_set for sid in request.stage_ids): diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 74234d8d6..c091ae5e8 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -4,12 +4,14 @@ import json import os import types -from dataclasses import fields, is_dataclass +from collections.abc import Mapping +from dataclasses import dataclass, fields, is_dataclass from typing import Any, get_args, get_origin from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm_omni.config.config_factory import with_trust_remote_code_override from vllm_omni.entrypoints.stage_utils import _to_dict from vllm_omni.inputs.data import OmniSamplingParams @@ -86,6 +88,64 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: return parsed +@dataclass(frozen=True) +class StageConfigInputs: + """Normalized inputs shared by all stage-config loading paths.""" + + model: str + kwargs: dict[str, Any] + trust_remote_code: bool | None + deploy_config_path: str | None + stage_overrides: dict[str, dict[str, Any]] | None + strategy_config_path: str | None + + +def prepare_stage_config_inputs( + model: str, + kwargs: Mapping[str, Any], + *, + trust_remote_code: bool | None, + snapshot_model: bool = False, +) -> StageConfigInputs: + """Normalize model/config arguments before resolving stage configs. + + The standard API worker, multi-API parent, and headless entrypoint must + apply the same legacy-key filtering, trust-remote-code precedence, and + config-file extraction before calling :func:`resolve_omni_config`. + """ + resolved_model = model + if snapshot_model: + # Keep model materialization at the CLI boundary for the parent-owned + # launch path. The API worker's OmniBase has already done this work. + from vllm_omni.entrypoints.omni_base import omni_snapshot_download + + resolved_model = omni_snapshot_download(model) + + resolved_kwargs = dict(kwargs) + for legacy_arg in ("stage_configs_path", "stage_configs"): + if legacy_arg in resolved_kwargs: + raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") + # Frontend process count is not a stage engine override. + resolved_kwargs.pop("api_server_count", None) + resolved_kwargs.pop("disable_log_stats", None) + if resolved_kwargs.get("diffusion_streaming_output") and resolved_kwargs.get("streaming_output") is None: + resolved_kwargs["streaming_output"] = True + resolved_kwargs.pop("model_tag", None) + + resolved_kwargs = with_trust_remote_code_override(resolved_kwargs, trust_remote_code) + deploy_config_path = resolved_kwargs.pop("deploy_config", None) + strategy_config_path = resolved_kwargs.pop("strategy_config", None) + stage_overrides = parse_stage_overrides(resolved_kwargs.pop("stage_overrides", None)) + return StageConfigInputs( + model=resolved_model, + kwargs=resolved_kwargs, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) + + def get_final_stage_id_for_e2e( output_modalities: list[str] | None, default_modalities: list[str], stage_list: list ) -> int: diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py index 64849a5aa..8e5f5832a 100644 --- a/vllm_omni/utils/tracking_parser.py +++ b/vllm_omni/utils/tracking_parser.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import argparse from typing import Any @@ -52,12 +52,28 @@ class TrackingNamespace(argparse.Namespace): else: setattr(self.unfiltered_ns, name, value) + def __getstate__(self) -> dict[str, Any]: + """Preserve wrapper state when API workers use ``spawn``. + + Our __dict__ property exposes the inner namespace, so default state + serialization would lose unfiltered_ns and explicit_keys. + """ + return { + "unfiltered_ns": self.unfiltered_ns, + "explicit_keys": self.explicit_keys, + } + + def __setstate__(self, state: dict[str, Any]) -> None: + object.__setattr__(self, "unfiltered_ns", state["unfiltered_ns"]) + object.__setattr__(self, "explicit_keys", state["explicit_keys"]) + def get_explicit_kwargs_dict(self): """Return a dict containing only the explicitly passed key-value pairs.""" return {k: v for k, v in vars(self.unfiltered_ns).items() if k in self.explicit_keys} def __getattr__(self, name: str) -> Any: - return getattr(self.unfiltered_ns, name) + namespace = object.__getattribute__(self, "unfiltered_ns") + return getattr(namespace, name) @property def __dict__(self): diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index e43c6f608..4d80d3f8d 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -812,7 +812,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] - elif self.uses_xdrope_dim > 0: + elif getattr(self, "uses_xdrope_dim", 0) > 0: positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: positions = self.positions[:num_tokens_padded] diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d3ee8de02..37d2ac122 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -690,7 +690,7 @@ class OmniGPUModelRunner(GPUModelRunner): self._init_mrope_positions(req_state) # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) - if self.uses_xdrope_dim > 0: + if getattr(self, "uses_xdrope_dim", 0) > 0: self._init_xdrope_positions(req_state) reqs_to_add.append(self.requests[req_id]) @@ -1164,7 +1164,7 @@ class OmniGPUModelRunner(GPUModelRunner): if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] - elif self.uses_xdrope_dim > 0: + elif getattr(self, "uses_xdrope_dim", 0) > 0: positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: positions = self.positions[:num_tokens_padded] @@ -1701,7 +1701,7 @@ class OmniGPUModelRunner(GPUModelRunner): if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_input_tokens] - elif self.uses_xdrope_dim > 0: + elif getattr(self, "uses_xdrope_dim", 0) > 0: positions = self.xdrope_positions.gpu[:, :num_input_tokens] else: positions = self.positions[:num_input_tokens] From 224152566a705dd0347e4b4c3c3e74ce7e19ed8f Mon Sep 17 00:00:00 2001 From: LOGO127 Date: Thu, 17 Sep 2026 14:04:37 +0800 Subject: [PATCH 23/99] [Model][PersonaPlex] Add paced multi-session validation (#7428) Signed-off-by: luozijian --- docs/design/fullduplex-personaplex.md | 108 ++++ .../personaplex_realtime_duplex.py | 567 ++++++++++++++--- .../test_personaplex_load_driver.py | 568 ++++++++++++++++++ 3 files changed, 1155 insertions(+), 88 deletions(-) create mode 100644 tests/e2e/online_serving/test_personaplex_load_driver.py diff --git a/docs/design/fullduplex-personaplex.md b/docs/design/fullduplex-personaplex.md index 33d125c97..6d32cd8e3 100644 --- a/docs/design/fullduplex-personaplex.md +++ b/docs/design/fullduplex-personaplex.md @@ -271,3 +271,111 @@ The required evidence is: Audio that is empty, all zero, non-finite, or only a protocol `listen` event is not a successful end-to-end result. + +### Paced multi-session load driver + +The existing driver keeps its two-session admission and slot-reuse validation +when `--sessions` is omitted. Passing `--sessions N` selects a separate, +bounded load mode; it does not replace the lifecycle test or change server +admission capacity. Start an existing PersonaPlex server configured to admit +the requested number of sessions, then run from a source checkout with the +vLLM-Omni client dependencies installed: + +```bash +python -m tests.e2e.online_serving.personaplex_realtime_duplex \ + --model nvidia/personaplex-7b-v1 \ + --input-wav /path/to/input.wav \ + --sessions 2 --load-frames 1000 --drain-s 2 \ + --server-revision SERVER_COMMIT --server-hardware A100-80GB \ + --output-dir /tmp/personaplex-load-n2 +``` + +Use a new output directory for every run. The driver refuses to overwrite an +existing `load-result.json`. The server labels are supplied by the operator, +not detected or verified by the client. Record the server command, package +versions, checkpoint revision and hardware alongside a published measurement. + +All attempted sessions finish admission before admitted sessions share a +common input start. A rejected session remains a failed row in the report; +it is not removed from the denominator. Each session receives at most +`--load-frames` input frames, without looping the WAV, and uses 80 ms pacing. +A stalled send cannot cause a catch-up burst. The timeline records send +schedule lateness so an overloaded client is visible rather than mistaken +for a faster server. This mode measures steady concurrency, not join/leave +interference; the original lifecycle mode still checks slot reuse. +Load mode does not probe an additional N+1 session for `resource_exhausted`; +the lifecycle mode owns that admission-limit check. A per-session +`client_pacing_warning` and the summary's `client_pacing_warning_sessions` +identify send-lateness p99 above one nominal frame (80 ms). This is a warning, +not an extra acceptance gate: read it alongside server-facing measurements +so a slow client is not mistaken for a slow server. + +The fixed drain window is part of the measurement configuration. Audio +arriving during cleanup is excluded, so closing cannot rescue a session that +missed its observation deadline. Cleanup is attempted for admitted sessions +even when audio validation fails. Protocol errors, missing/invalid audio, +shared response IDs and cleanup failures make the report fail and the CLI +exit nonzero. Cancellation closes transports and propagates to the caller. +`--cleanup-timeout-s` (default 5 s) bounds session-close send and acknowledgement +together, independently of `--timeout-s`. The subsequent WebSocket close +handshake has the same separate bound; an unresponsive close acknowledgement +cannot hold cancellation for the full inference timeout. + +`load-result.json` retains per-session send and audio-packet timelines using +the client's monotonic clock. It does not export handshake credentials or +raw audio payloads. Metric definitions are explicit: + +Both modes validate each nonempty audio packet's strict base64 encoding, PCM16 +sample alignment, 24 kHz sample rate and response identity. Valid aggregate +audio cannot hide a malformed packet. In load mode, rejected packets make the session fail +and appear in `invalid_audio_packets` with their event index, receive time and +fixed diagnostic code; they do not contribute samples or timing intervals. +Empty audio deltas remain legal and contribute no samples. This check does +not change the public client's decoding. The lifecycle scenario and dispatch +remain the same, but malformed packets that aggregate checks once tolerated +now fail that mode as well. + +The probe accepts both the current OpenAI Realtime audio event name +`response.output_audio.delta` and the legacy `response.audio.delta`. Validation +uses the wire payload directly rather than relying on a particular client +library's alias table, so a server-side event-name migration cannot silently +turn real audio into a zero-output measurement. + +| Field | Meaning | +| --- | --- | +| `client_audio_packet_interval_ms` | Consecutive nonempty packet arrivals: median, linear p99, max, and zero-based index of the first maximal interval. Interval index 0 ends at packet 1. One packet can contain multiple codec frames; this is not server tick latency. | +| `client_first_audio_after_stream_start_ms` | First audio receipt minus first input send start; admission is excluded. | +| `client_stream_rtf` | First input send start to last audio receipt, divided by received audio duration. This includes real-time input pacing and is not an inference-only RTF. | +| `client_send_lateness_ms` | Send start minus its original shared schedule, clipped at zero. | +| `frame_deficit` | Integer input-frame minus output-frame count from the shared session validation; missing when valid whole-frame accounting is unavailable. | +| `output_frames`, `voiced_frames`, `silent_frames`, `frame_coverage_ratio`, `audio_rms`, `audio_chunks` | Session-validation statistics retained on success and on acceptance failure once computed. | +| `acceptance_check` | Failed local check, with numeric diagnostics in `error`; other exceptions retain phase/type only, without remote handshake or metadata text. | +| `audio_underrun_s`, `audio_underrun_event_count`, `audio_continuity_ok` | Shared `compute_continuity_stats` on valid PCM16 packet arrivals, with its default 100 ms underrun threshold; absent audio gives missing values rather than a continuity pass. | + +No output gives missing latency/RTF values, not zero. By default RTF is +diagnostic; `--max-client-rtf` adds an explicit ceiling. The existing frame +coverage, audible-frame and scheduler-data-plane checks still apply. A longer +`--drain-s` can demonstrate eventual completion under backlog, but it must not +be reported as real-time capacity when the fixed-window run failed or client +RTF exceeded the stated target. Likewise, the first live session can pay +one-time model/codec/JIT costs; either use the server's configured duplex +warmup path or label that run as cold-start evidence before comparing steady +capacity. Passing client-only metrics does not certify model quality or a +real-time capacity target on another GPU. Begin with N=1 and N=2; higher N is a +capacity experiment, not a promised property of an A100 or of this patch. +Continuity is a client-arrival diagnostic, not proof of speaker playback and +not an additional pass/fail gate. Packet-interval and send-lateness summaries +retain their documented linear p99; the older client interval helper reports +rounded nearest-rank p95 instead and is not substituted for that schema. + +Deterministic driver and localhost WebSocket checks need no model or GPU: + +```bash +python -m pytest tests/e2e/online_serving/test_personaplex_load_driver.py \ + -m 'core_model and cpu' --run-level core_model -q +``` + +These checks test the driver against controlled protocol responses. Run the +load-driver command against real weights separately before reporting model-serving +performance. Server tick timestamps and generic TTFT/TPOT aggregation remain +the responsibility of the existing duplex metrics paths, not this driver. diff --git a/tests/e2e/online_serving/personaplex_realtime_duplex.py b/tests/e2e/online_serving/personaplex_realtime_duplex.py index 20da1598d..78f012548 100644 --- a/tests/e2e/online_serving/personaplex_realtime_duplex.py +++ b/tests/e2e/online_serving/personaplex_realtime_duplex.py @@ -10,24 +10,24 @@ import base64 import hashlib import json import math +import time import uuid import wave -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence from pathlib import Path +from typing import Protocol from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import numpy as np try: import websockets - from websockets.exceptions import ConnectionClosed + from websockets.exceptions import ConnectionClosed, WebSocketException except ImportError as exc: # pragma: no cover - driver dependency. raise SystemExit("Install websockets first: pip install websockets") from exc -from vllm_omni.clients.duplex import ( - EventCollector, - write_pcm16_wav, -) +from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats +from vllm_omni.clients.duplex import EventCollector, write_pcm16_wav from vllm_omni.clients.duplex import ( wait_for_condition as wait_for, ) @@ -35,6 +35,24 @@ from vllm_omni.clients.duplex import ( SAMPLE_RATE_HZ = 24_000 FRAME_SAMPLES = 1_920 FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ +AUDIO_DELTA_EVENT_TYPES = frozenset({"response.audio.delta", "response.output_audio.delta"}) + + +class _AudioAcceptanceError(AssertionError): + """A local acceptance failure with numeric diagnostics safe to report.""" + + def __init__(self, check: str, message: str, stats: dict[str, object]) -> None: + super().__init__(message) + self.check = check + self.stats = dict(stats) + + +class _ProbeTransport(Protocol): + async def recv(self) -> str | bytes: ... + + async def send(self, message: str) -> None: ... + + async def close(self) -> None: ... class RawRealtimeProbe: @@ -46,29 +64,33 @@ class RawRealtimeProbe: keeps a transport-only client and reads everything from ``.events``. """ - def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None: + def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024, close_timeout_s: float = 10.0) -> None: self.url = url self.max_size = max_size + self.close_timeout_s = close_timeout_s self.events = EventCollector() - self._ws = None + self._ws: _ProbeTransport | None = None self._reader_task: asyncio.Task[None] | None = None async def __aenter__(self) -> RawRealtimeProbe: - self._ws = await websockets.connect(self.url, max_size=self.max_size) + self._ws = await websockets.connect(self.url, max_size=self.max_size, close_timeout=self.close_timeout_s) self._reader_task = asyncio.create_task(self._read_events()) return self async def __aexit__(self, exc_type, exc, traceback) -> None: - if self._ws is not None: - await self._ws.close() - if self._reader_task is not None: - self._reader_task.cancel() - try: - await self._reader_task - except asyncio.CancelledError: - pass + try: + if self._ws is not None: + await self._ws.close() + finally: + if self._reader_task is not None: + self._reader_task.cancel() + try: + await self._reader_task + except asyncio.CancelledError: + pass async def _read_events(self) -> None: + assert self._ws is not None try: while True: raw = await self._ws.recv() @@ -81,6 +103,7 @@ class RawRealtimeProbe: return async def send(self, event: dict[str, object]) -> None: + assert self._ws is not None await self._ws.send(json.dumps(event)) @@ -142,6 +165,38 @@ def _events(client: RawRealtimeProbe, event_type: str) -> list[dict[str, object] return [event for event in client.events.events if event.get("type") == event_type] +def _audio_events(client: RawRealtimeProbe) -> list[dict[str, object]]: + """Return audio deltas from both current and legacy Realtime event names.""" + return [event for event in client.events.events if event.get("type") in AUDIO_DELTA_EVENT_TYPES] + + +def _validated_audio_event(event: dict[str, object]) -> bytes: + """Decode one wire audio delta without depending on client event aliases.""" + encoded = event.get("delta") + if not isinstance(encoded, str): + raise ValueError("missing_audio_delta") + try: + chunk = base64.b64decode(encoded, validate=True) + except ValueError: + raise ValueError("invalid_audio_base64") from None + if chunk: + if len(chunk) % 2: + raise ValueError("unaligned_pcm16_packet") + if event.get("sample_rate_hz") != SAMPLE_RATE_HZ: + raise ValueError("invalid_audio_sample_rate") + if not isinstance(event.get("response_id"), str) or not event.get("response_id"): + raise ValueError("missing_audio_response_id") + return chunk + + +def _validated_audio_chunks(client: RawRealtimeProbe) -> list[tuple[dict[str, object], bytes]]: + return [(event, _validated_audio_event(event)) for event in _audio_events(client)] + + +def _audio_bytes(client: RawRealtimeProbe) -> bytes: + return b"".join(chunk for _, chunk in _validated_audio_chunks(client) if chunk) + + async def _open_session( args: argparse.Namespace, *, @@ -151,22 +206,7 @@ async def _open_session( ) -> tuple[RawRealtimeProbe, dict[str, object]]: client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id)) await client.__aenter__() - await client.send( - { - "type": "session.update", - "session": { - "session_id": session_id, - "model": args.model, - "modalities": ["audio", "text"], - "input_audio_format": "pcm_f32le", - "output_audio_format": "pcm16", - "voice": args.voice, - "instructions": persona, - "turn_detection": None, - "extra_body": {"auto_response": True}, - }, - } - ) + await client.send(_session_update(args, session_id=session_id, persona=persona)) event_type = "error" if expect_error else "session.created" await wait_for( lambda: client.events.count(event_type) > 0, @@ -176,6 +216,23 @@ async def _open_session( return client, _events(client, event_type)[-1] +def _session_update(args: argparse.Namespace, *, session_id: str, persona: str) -> dict[str, object]: + return { + "type": "session.update", + "session": { + "session_id": session_id, + "model": args.model, + "modalities": ["audio", "text"], + "input_audio_format": "pcm_f32le", + "output_audio_format": "pcm16", + "voice": args.voice, + "instructions": persona, + "turn_detection": None, + "extra_body": {"auto_response": True}, + }, + } + + async def _close_session(client: RawRealtimeProbe, *, timeout_s: float) -> None: await client.send({"type": "session.close"}) await wait_for( @@ -200,23 +257,24 @@ async def _stream_frames( frame_count = min(frame_count, max_frames) for seq in range(frame_count): frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES] - frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)) - frame = np.ascontiguousarray(frame, dtype="= 1e-4), - } - ) + await client.send(_frame_event(frame, seq)) await asyncio.sleep(FRAME_PERIOD_S) return frame_count +def _frame_event(frame: np.ndarray, seq: int) -> dict[str, object]: + frame = np.ascontiguousarray(np.pad(frame, (0, FRAME_SAMPLES - frame.size)), dtype="= 1e-4), + } + + def _audio_frame_stats( raw: bytes, *, @@ -242,11 +300,7 @@ def _audio_frame_stats( def _response_ids(client: RawRealtimeProbe) -> set[str]: - return { - str(event["response_id"]) - for event in _events(client, "response.output_audio.delta") - if isinstance(event.get("response_id"), str) - } + return {str(event["response_id"]) for event in _audio_events(client) if isinstance(event.get("response_id"), str)} def _session_result( @@ -256,40 +310,66 @@ def _session_result( args: argparse.Namespace, minimum_chunks: int, ) -> tuple[bytes, set[str], dict[str, object]]: - chunks = sum(len(items) for items in client.events.response_audio.values()) - raw = client.events.audio_bytes() - if chunks < minimum_chunks or not raw or len(raw) % 2: - raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}") + if isinstance(client, RawRealtimeProbe): + audio_chunks = _validated_audio_chunks(client) + chunks = sum(bool(chunk) for _, chunk in audio_chunks) + raw = b"".join(chunk for _, chunk in audio_chunks if chunk) + else: + # Unit fixtures may provide an already-decoded EventCollector without + # retaining the base64 wire field. Real websocket probes always take + # the strict wire-validation path above. + chunks = sum(len(items) for items in client.events.response_audio.values()) + raw = client.events.audio_bytes() + stats: dict[str, object] = {"audio_chunks": chunks, "audio_samples": len(raw) // 2} + if not raw or len(raw) % 2: + raise _AudioAcceptanceError("pcm16_output", f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}", stats) pcm = np.frombuffer(raw, dtype=" async def run(args: argparse.Namespace) -> dict[str, object]: + if getattr(args, "sessions", None) is not None: + return await run_load(args) input_path = Path(args.input_wav) input_identity = _input_identity( input_path, @@ -381,14 +464,12 @@ async def run(args: argparse.Namespace) -> dict[str, object]: ) await asyncio.gather( wait_for( - lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit, + lambda: len(_audio_bytes(primary)) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit, timeout_s=args.timeout_s, label="primary frame coverage", ), wait_for( - lambda: ( - len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit - ), + lambda: (len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit), timeout_s=args.timeout_s, label="secondary frame coverage", ), @@ -426,16 +507,14 @@ async def run(args: argparse.Namespace) -> dict[str, object]: await asyncio.gather( wait_for( lambda: ( - len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) - >= secondary_total_frames - args.max_frame_deficit + len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_total_frames - args.max_frame_deficit ), timeout_s=args.timeout_s, label="survivor frame coverage after slot recycle", ), wait_for( lambda: ( - len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES) - >= replacement_frames - args.max_frame_deficit + len(_audio_bytes(replacement)) // (2 * FRAME_SAMPLES) >= replacement_frames - args.max_frame_deficit ), timeout_s=args.timeout_s, label="replacement frame coverage", @@ -493,6 +572,284 @@ async def run(args: argparse.Namespace) -> dict[str, object]: return result +def _distribution(values: list[float]) -> dict[str, int | float | None]: + if not values: + return {"count": 0, "median": None, "p99": None, "max": None, "argmax_index": None} + return { + "count": len(values), + "median": float(np.median(values)), + "p99": float(np.quantile(values, 0.99)), + "max": max(values), + "argmax_index": int(np.argmax(values)), + } + + +def _load_metrics(client: RawRealtimeProbe, sends: list[tuple[float, float, float]]) -> dict[str, object]: + # Use the collector's local clock, never a timestamp supplied by the server. + packets: list[dict[str, object]] = [] + response_ids: set[str] = set() + packet_times: list[float] = [] + packet_bytes: list[int] = [] + invalid_packets: list[dict[str, object]] = [] + samples = 0 + for index, (raw, received) in enumerate(zip(client.events.events, client.events.event_received_at_s, strict=True)): + if raw.get("type") not in AUDIO_DELTA_EVENT_TYPES: + continue + try: + chunk = _validated_audio_event(raw) + except ValueError as exc: + # Only fixed diagnostic codes from the validator are exported, + # never audio payloads or details from the remote endpoint. + invalid_packets.append({"event_index": index, "received_at_s": received, "code": str(exc)}) + continue + if chunk: + samples += len(chunk) // 2 + packet_times.append(received) + packet_bytes.append(len(chunk)) + response_id = raw.get("response_id") + packets.append({"received_at_s": received, "samples": len(chunk) // 2, "response_id": response_id}) + if isinstance(response_id, str) and response_id: + response_ids.add(response_id) + intervals = [(b - a) * 1000 for a, b in zip(packet_times, packet_times[1:])] + first_send = sends[0][1] if sends else None + audio_duration = samples / SAMPLE_RATE_HZ + first_audio_ms = (packet_times[0] - first_send) * 1000 if packet_times and first_send is not None else None + client_rtf = (packet_times[-1] - first_send) / audio_duration if audio_duration and first_send is not None else None + continuity = compute_continuity_stats(packet_times, packet_bytes, SAMPLE_RATE_HZ) if packet_times else None + lateness = _distribution([max(0.0, sent - planned) * 1000 for planned, sent, _ in sends]) + p99_lateness = lateness["p99"] + return { + "input_frames": len(sends), + "output_samples": samples, + # Populate frame accounting from _session_result, including on failure. + "frame_deficit": None, + "response_ids": sorted(response_ids), + "client_first_audio_after_stream_start_ms": first_audio_ms, + "client_stream_rtf": client_rtf, + "client_audio_packet_interval_ms": _distribution(intervals), + "client_send_lateness_ms": lateness, + "client_pacing_warning": isinstance(p99_lateness, float) and p99_lateness > FRAME_PERIOD_S * 1000, + "audio_underrun_s": continuity.max_underrun_s if continuity else None, + "audio_underrun_event_count": continuity.underrun_event_count if continuity else None, + "audio_continuity_ok": continuity.is_continuous if continuity else None, + "input_timeline": [ + {"frame_index": i, "planned_at_s": planned, "send_started_at_s": sent, "send_completed_at_s": completed} + for i, (planned, sent, completed) in enumerate(sends) + ], + "audio_packet_timeline": packets, + "invalid_audio_packets": invalid_packets, + } + + +def _check_load_connection(client: RawRealtimeProbe) -> None: + if client.events.errors(): + raise RuntimeError("server emitted a Realtime error (see the session error count)") + if client._reader_task is not None and client._reader_task.done(): + client._reader_task.result() # Propagate malformed JSON/reader failures. + raise RuntimeError("connection closed before session completion") + + +async def _wait_load_event(client: RawRealtimeProbe, event_type: str, timeout_s: float) -> None: + def received() -> bool: + if client.events.count(event_type): + return True + _check_load_connection(client) + return False + + await wait_for(received, timeout_s=timeout_s, label=event_type) + + +async def _paced_load_frames( + client: RawRealtimeProbe, + pcm: np.ndarray, + *, + epoch: float, + timeout_s: float, + sends: list[tuple[float, float, float]], + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> None: + next_send = epoch + for seq in range(math.ceil(pcm.size / FRAME_SAMPLES)): + await sleep(max(0.0, next_send - clock())) + _check_load_connection(client) + event = _frame_event(pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES], seq) + started = clock() + await asyncio.wait_for(client.send(event), timeout=timeout_s) + completed = clock() + sends.append((epoch + seq * FRAME_PERIOD_S, started, completed)) + # Do not burst to catch up after a stalled send; report accumulated lag. + next_send = max(epoch + (seq + 1) * FRAME_PERIOD_S, completed + FRAME_PERIOD_S) + + +async def _request_load_close(client: RawRealtimeProbe, timeout_s: float) -> None: + await client.send({"type": "session.close"}) + await _wait_load_event(client, "session.closed", timeout_s) + + +async def _run_load_session( + args: argparse.Namespace, + pcm: np.ndarray, + index: int, + ready: asyncio.Future[None], + start: asyncio.Future[float], +) -> dict[str, object]: + session_id = f"personaplex-load-{index}-{uuid.uuid4().hex}" + client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id), close_timeout_s=args.cleanup_timeout_s) + sends: list[tuple[float, float, float]] = [] + error: str | None = None + cleanup_error: str | None = None + metrics: dict[str, object] | None = None + session_stats: dict[str, object] = {} + acceptance_check: str | None = None + phase = "admission" + try: + async with client: + try: + await asyncio.wait_for( + client.send(_session_update(args, session_id=session_id, persona=args.persona)), args.timeout_s + ) + await _wait_load_event(client, "session.created", args.timeout_s) + capabilities = _capabilities(_events(client, "session.created")[-1]) + if ( + capabilities.get("chunk_period_ms") != 80 + or capabilities.get("supports_multi_session_same_replica") is not True + ): + raise AssertionError("server did not advertise PersonaPlex multi-session capabilities") + ready.set_result(None) + phase = "stream" + epoch = await asyncio.shield(start) + await _paced_load_frames(client, pcm, epoch=epoch, timeout_s=args.timeout_s, sends=sends) + phase = "drain" + # Fixed, reported observation window. Closing must not extend + # it and let a slow session pass by flushing after the cutoff. + await asyncio.sleep(args.drain_s) + _check_load_connection(client) + _, _, session_stats = _session_result( + client, input_frames=len(sends), args=args, minimum_chunks=args.minimum_audio_chunks + ) + finally: + metrics = _load_metrics(client, sends) + if client.events.count("session.created") and not client.events.count("session.closed"): + try: + # Send and acknowledgement share one short cleanup budget, + # independent of the potentially long inference timeout. + await asyncio.wait_for( + _request_load_close(client, args.cleanup_timeout_s), args.cleanup_timeout_s + ) + except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, WebSocketException) as exc: + cleanup_error = f"cleanup: {type(exc).__name__}" + except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, AssertionError, WebSocketException) as exc: + if isinstance(exc, _AudioAcceptanceError): + acceptance_check = exc.check + session_stats = exc.stats + error = f"{phase}: AssertionError: {exc}" + else: + # Other assertions/errors may embed handshake credentials or wire + # metadata. Only our local numeric acceptance diagnostics are safe. + error = f"{phase}: {type(exc).__name__}" + finally: + # Admission failure must not leave the other sessions at the barrier. + if not ready.done(): + ready.set_result(None) + if metrics is None: + metrics = _load_metrics(client, sends) + if cleanup_error: + error = f"{error}; {cleanup_error}" if error else cleanup_error + if client.events.errors(): + error = error or "server emitted a Realtime error" + if metrics["invalid_audio_packets"]: + error = error or "invalid audio packets" + first_audio = metrics["client_first_audio_after_stream_start_ms"] + if isinstance(first_audio, float) and first_audio < 0: + error = error or "audio arrived before the first input frame" + rtf = metrics["client_stream_rtf"] + if args.max_client_rtf is not None and isinstance(rtf, float) and rtf > args.max_client_rtf: + error = error or f"client_stream_rtf exceeds {args.max_client_rtf}" + return { + "index": index, + "session_id": session_id, + "ok": error is None, + "error": error, + "acceptance_check": acceptance_check, + "server_error_count": len(client.events.errors()), + "server_error_codes": [ + str(body.get("code", "unknown")) if isinstance(body := event.get("error"), dict) else "unknown" + for event in client.events.errors() + ], + **metrics, + **session_stats, + } + + +async def run_load(args: argparse.Namespace) -> dict[str, object]: + input_path = Path(args.input_wav) + identity = _input_identity(input_path, expected_sha256=args.expected_input_sha256) + pcm = _read_wav_as_float32(input_path) + if args.tail_s: + pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)]) + pcm = pcm[: args.load_frames * FRAME_SAMPLES] + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + # Refuse to silently overwrite evidence from a previous sweep. + result_path = output_dir / "load-result.json" + if result_path.exists(): + raise FileExistsError(result_path) + loop = asyncio.get_running_loop() + start: asyncio.Future[float] = loop.create_future() + ready: list[asyncio.Future[None]] = [loop.create_future() for _ in range(args.sessions)] + tasks = [asyncio.create_task(_run_load_session(args, pcm, i, r, start)) for i, r in enumerate(ready)] + try: + await asyncio.gather(*ready) + start.set_result(time.monotonic()) + rows = await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + seen: set[str] = set() + shared: set[str] = set() + for row in rows: + assert isinstance(row["response_ids"], list) + response_ids = set(row["response_ids"]) + shared.update(seen & response_ids) + seen.update(response_ids) + for row in rows: + assert isinstance(row["response_ids"], list) + if shared.intersection(row["response_ids"]): + row.update(ok=False, error="concurrent sessions shared a response id") + result = { + "schema_version": 1, + "mode": "paced_multi_session", + "model": args.model, + "requested_sessions": args.sessions, + "passed_sessions": sum(row["ok"] is True for row in rows), + "ok": all(row["ok"] is True for row in rows), + "input": {**identity, "tail_s": args.tail_s, "load_frames": args.load_frames}, + "frame_period_s": FRAME_PERIOD_S, + "drain_s": args.drain_s, + "cleanup_timeout_s": args.cleanup_timeout_s, + "client_pacing_warning_threshold_ms": FRAME_PERIOD_S * 1000, + "client_pacing_warning_sessions": [row["index"] for row in rows if row["client_pacing_warning"]], + "max_client_rtf": args.max_client_rtf, + "acceptance": { + "max_frame_deficit": args.max_frame_deficit, + "min_voiced_frames": args.min_voiced_frames, + "minimum_audio_chunks": args.minimum_audio_chunks, + "voiced_frame_rms_threshold": args.voiced_frame_rms_threshold, + }, + "shared_response_ids": sorted(shared), + "sessions": rows, + "reported_server": {"revision": args.server_revision, "hardware": args.server_hardware}, + "measurement_origin": "client monotonic clock; admission excluded; streaming and drain retained", + "interval_definition": "nonempty audio packet arrivals, not server ticks or per-frame inference latency", + } + with result_path.open("x", encoding="utf-8") as stream: + json.dump(result, stream, ensure_ascii=False, indent=2, allow_nan=False) + return result + + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1") @@ -510,6 +867,12 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--tail-s", type=float, default=2.0) parser.add_argument("--drain-s", type=float, default=2.0) parser.add_argument("--timeout-s", type=float, default=180.0) + parser.add_argument( + "--cleanup-timeout-s", + type=float, + default=5.0, + help="Load mode: shared session-close/ack budget; also bounds the websocket close handshake.", + ) parser.add_argument( "--max-frame-deficit", type=int, @@ -536,11 +899,39 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=0, help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.", ) - return parser.parse_args(argv) + parser.add_argument( + "--sessions", type=int, help="Opt into paced load mode with N sessions; omit for the lifecycle test." + ) + parser.add_argument( + "--load-frames", type=int, default=1000, help="Maximum input frames per load session (no repetition)." + ) + parser.add_argument("--max-client-rtf", type=float, help="Optional client stream-window RTF acceptance ceiling.") + parser.add_argument("--server-revision", help="User-reported server commit/version; not verified by this client.") + parser.add_argument( + "--server-hardware", help="User-reported server hardware; not inferred from the client machine." + ) + args = parser.parse_args(argv) + if args.sessions is not None: + if not 1 <= args.sessions <= 64 or not 1 <= args.load_frames <= 10000: + parser.error("load mode requires sessions in [1,64] and load-frames in [1,10000]") + for name in ("timeout_s", "cleanup_timeout_s", "drain_s", "tail_s"): + value = getattr(args, name) + if not math.isfinite(value) or value < 0 or (name in ("timeout_s", "cleanup_timeout_s") and value == 0): + parser.error(f"{name} must be finite and non-negative (timeout must be positive)") + if args.max_client_rtf is not None and (not math.isfinite(args.max_client_rtf) or args.max_client_rtf <= 0): + parser.error("max-client-rtf must be finite and positive") + if args.max_frame_deficit < 0 or args.min_voiced_frames < 1 or args.minimum_audio_chunks < 1: + parser.error("frame-deficit must be non-negative; voiced frames and audio chunks must be positive") + if not math.isfinite(args.voiced_frame_rms_threshold) or args.voiced_frame_rms_threshold < 0: + parser.error("voiced-frame-rms-threshold must be finite and non-negative") + return args def main() -> None: - print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2)) + result = asyncio.run(run(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False)) + if result.get("mode") == "paced_multi_session" and not result["ok"]: + raise SystemExit(1) if __name__ == "__main__": diff --git a/tests/e2e/online_serving/test_personaplex_load_driver.py b/tests/e2e/online_serving/test_personaplex_load_driver.py new file mode 100644 index 000000000..8e9a6a28d --- /dev/null +++ b/tests/e2e/online_serving/test_personaplex_load_driver.py @@ -0,0 +1,568 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import base64 +import json +import wave +from contextlib import asynccontextmanager +from typing import TypedDict + +import numpy as np +import pytest +import websockets + +from tests.e2e.online_serving import personaplex_realtime_duplex as driver + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _ServerState(TypedDict): + index: int + closed: bool + frames: int + close_requested: bool + + +def _args(tmp_path, *extra): + wav_path = tmp_path / "input.wav" + with wave.open(str(wav_path), "wb") as stream: + stream.setnchannels(1) + stream.setsampwidth(2) + stream.setframerate(driver.SAMPLE_RATE_HZ) + stream.writeframes(np.full(driver.FRAME_SAMPLES * 2, 2000, dtype="= 1 + assert report["client_pacing_warning"] is False + assert "secret" not in json.dumps(report) + assert "-999" not in json.dumps(report) + + +def test_empty_output_has_missing_not_zero_latency(): + report = driver._load_metrics(driver.RawRealtimeProbe("ws://unused"), []) + assert report["client_stream_rtf"] is None + assert report["audio_continuity_ok"] is None + assert report["audio_underrun_s"] is None + intervals = report["client_audio_packet_interval_ms"] + assert isinstance(intervals, dict) + assert intervals["p99"] is None + + +def test_failed_acceptance_retains_frame_counts_without_remote_data(): + client = driver.RawRealtimeProbe("ws://unused") + client.events.add(_audio(), received_at_s=1.0) + args = driver.parse_args( + ["--model", "unused", "--input-wav", "unused.wav", "--max-frame-deficit", "0", "--min-voiced-frames", "1"] + ) + with pytest.raises(AssertionError) as raised: + driver._session_result(client, input_frames=2, args=args, minimum_chunks=1) + assert getattr(raised.value, "check", None) == "frame_deficit" + stats = getattr(raised.value, "stats", {}) + assert stats["frame_deficit"] == 1 and type(stats["frame_deficit"]) is int + assert stats["output_frames"] == stats["voiced_frames"] == 1 + assert stats["frame_coverage_ratio"] == 0.5 + + +@pytest.mark.parametrize("event_type", sorted(driver.AUDIO_DELTA_EVENT_TYPES)) +def test_session_result_accepts_current_and_legacy_audio_event_names(tmp_path, event_type): + client = driver.RawRealtimeProbe("ws://unused") + client.events.add(_audio(samples=driver.FRAME_SAMPLES, event_type=event_type), received_at_s=1.0) + raw, response_ids, stats = driver._session_result( + client, + input_frames=1, + args=_args(tmp_path, "--sessions", "1"), + minimum_chunks=1, + ) + assert len(raw) == driver.FRAME_SAMPLES * 2 + assert response_ids == {"r"} + assert stats["frame_deficit"] == 0 + + +def test_legacy_audio_event_name_contributes_load_metrics(): + client = driver.RawRealtimeProbe("ws://unused") + client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5, event_type="response.audio.delta"), received_at_s=10.4) + report = driver._load_metrics(client, [(10.0, 10.0, 10.01)]) + assert report["output_samples"] == driver.FRAME_SAMPLES * 5 + assert report["response_ids"] == ["r"] + assert report["invalid_audio_packets"] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sessions", [1, 2, 4]) +async def test_real_websocket_load_success(tmp_path, sessions): + args = _args(tmp_path, "--sessions", str(sessions)) + async with _server() as (url, states): + args.url = url + result = await driver.run(args) + assert result["ok"] is True + assert result["passed_sessions"] == sessions + rows = result["sessions"] + assert isinstance(rows, list) + assert len(rows) == sessions + assert all(row["frame_deficit"] == 0 for row in rows) + for row in rows: + assert type(row["frame_deficit"]) is int + assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 2 + assert row["silent_frames"] == 0 + assert row["frame_coverage_ratio"] == 1.0 + assert row["audio_rms"] == pytest.approx(2000 / 32768) + assert row["acceptance_check"] is None + assert isinstance(row["audio_continuity_ok"], bool) + assert all(s["closed"] and s["close_requested"] for s in states) + assert "must-not-be-exported" not in json.dumps(result) + saved = json.loads((tmp_path / "out/load-result.json").read_text()) + assert saved == result + origins = [row["input_timeline"][0]["planned_at_s"] for row in rows] + assert len(set(origins)) == 1 + + +@pytest.mark.asyncio +async def test_real_websocket_load_accepts_legacy_audio_event_name(tmp_path): + args = _args(tmp_path, "--sessions", "2") + async with _server("legacy_audio_event") as (url, states): + args.url = url + result = await driver.run(args) + assert result["ok"] is True + assert result["passed_sessions"] == 2 + assert result["shared_response_ids"] == [] + assert all(s["closed"] and s["close_requested"] for s in states) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode", + [ + "reject_second", + "silent", + "disconnect", + "no_close_ack", + "shared_response", + "late_on_close", + "error_on_close", + "bad_json", + "missing_capabilities", + ], +) +async def test_real_websocket_failures_remain_in_results(tmp_path, mode): + args = _args(tmp_path) + async with _server(mode) as (url, states): + args.url = url + result = await asyncio.wait_for(driver.run(args), timeout=4) + assert result["ok"] is False + assert "must-not-be-exported" not in json.dumps(result) + assert result["requested_sessions"] == 2 + rows = result["sessions"] + assert isinstance(rows, list) + assert len(rows) == 2 + assert all(s["closed"] for s in states) + if mode == "reject_second": + assert result["passed_sessions"] == 1 + rejected = next(row for row in rows if not row["ok"]) + assert rejected["input_frames"] == 0 + assert rejected["client_stream_rtf"] is None + assert rejected["server_error_codes"] == ["resource_exhausted"] + if mode == "silent": + # A failed measurement must still release its admitted session. + assert all(s["close_requested"] for s in states) + if mode == "late_on_close": + assert all(row["output_samples"] == 0 for row in rows) + if mode == "shared_response": + assert result["passed_sessions"] == 0 + + +@pytest.mark.asyncio +async def test_existing_report_is_not_overwritten(tmp_path): + args = _args(tmp_path) + path = tmp_path / "out/load-result.json" + path.parent.mkdir() + path.write_text("previous measurement") + with pytest.raises(FileExistsError): + await driver.run(args) + assert path.read_text() == "previous measurement" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["success", "no_close_ack"]) +async def test_cancellation_closes_all_connections(tmp_path, mode): + args = _args(tmp_path, "--drain-s", "10", "--timeout-s", "30") + async with _server(mode) as (url, states): + args.url = url + task = asyncio.create_task(driver.run(args)) + + async def started(): + while len(states) < 2 or not all(state["frames"] for state in states): + await asyncio.sleep(0.01) + + await asyncio.wait_for(started(), timeout=3) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + assert all(s["closed"] for s in states) + + +@pytest.mark.asyncio +async def test_client_rtf_ceiling_marks_slow_sessions_failed(tmp_path): + args = _args(tmp_path, "--max-client-rtf", "0.001") + async with _server() as (url, _): + args.url = url + result = await driver.run(args) + assert result["passed_sessions"] == 0 + rows = result["sessions"] + assert isinstance(rows, list) + assert all("client_stream_rtf exceeds" in row["error"] for row in rows) + + +@pytest.mark.asyncio +async def test_stalled_send_does_not_cause_catchup_burst(): + class Clock: + now = 10.0 + + def monotonic(self): + return self.now + + async def sleep(self, delay): + self.now += delay + + class Client(driver.RawRealtimeProbe): + async def send(self, event): + if event["audio_end_ms"] == 80: + clock.now += 0.25 + + clock = Clock() + sends: list[tuple[float, float, float]] = [] + await driver._paced_load_frames( + Client("ws://unused"), + np.zeros(driver.FRAME_SAMPLES * 3, dtype="= sends[0][2] + driver.FRAME_PERIOD_S + assert sends[1][1] - sends[1][0] == pytest.approx(0.25) + assert sends[2][1] >= sends[1][2] + driver.FRAME_PERIOD_S + metrics = driver._load_metrics(Client("ws://unused"), sends) + assert metrics["client_pacing_warning"] is True + assert asyncio.sleep != clock.sleep + + +@pytest.mark.asyncio +async def test_failed_frame_acceptance_preserves_numeric_diagnostics(tmp_path): + args = _args(tmp_path, "--sessions", "1") + async with _server("short_audio") as (url, states): + args.url = url + result = await driver.run(args) + rows = result["sessions"] + assert isinstance(rows, list) + row = rows[0] + assert row["ok"] is False + assert row["acceptance_check"] == "frame_deficit" + assert row["frame_deficit"] == 1 and type(row["frame_deficit"]) is int + assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 1 + assert row["frame_coverage_ratio"] == 0.5 + assert row["audio_rms"] == pytest.approx(2000 / 32768) + assert "input=2" in row["error"] and "maximum_deficit=0" in row["error"] + assert all(state["closed"] and state["close_requested"] for state in states) + assert json.loads((tmp_path / "out/load-result.json").read_text()) == result + + +@pytest.mark.asyncio +async def test_acceptance_failure_does_not_export_remote_metadata(tmp_path): + args = _args(tmp_path, "--sessions", "1") + async with _server("private_metadata") as (url, _): + args.url = url + result = await driver.run(args) + rows = result["sessions"] + assert isinstance(rows, list) + row = rows[0] + assert row["ok"] is False + assert row["acceptance_check"] == "scheduler_data_plane" + assert "remote-detail" not in json.dumps(result) + assert "must-not-be-exported" not in json.dumps(result) + + +@pytest.mark.parametrize("ok", [True, False]) +def test_load_main_exit_status(tmp_path, monkeypatch, ok): + args = _args(tmp_path) + + async def completed(_args): + return {"mode": "paced_multi_session", "ok": ok} + + monkeypatch.setattr(driver, "parse_args", lambda: args) + monkeypatch.setattr(driver, "run", completed) + if ok: + driver.main() + else: + with pytest.raises(SystemExit) as exc: + driver.main() + assert exc.value.code == 1 + + +@pytest.mark.asyncio +async def test_handshake_rejection_is_retained_without_server_details(tmp_path, monkeypatch): + from websockets.exceptions import InvalidHandshake + + async def rejected(_self): + raise InvalidHandshake("credentials-in-remote-handshake-detail") + + monkeypatch.setattr(driver.RawRealtimeProbe, "__aenter__", rejected) + result = await asyncio.wait_for(driver.run(_args(tmp_path)), timeout=3) + assert result["ok"] is False + rows = result["sessions"] + assert isinstance(rows, list) and len(rows) == 2 + assert all(row["error"] == "admission: InvalidHandshake" for row in rows) + assert "credentials-in-remote-handshake-detail" not in json.dumps(result) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sessions", [1, 2]) +@pytest.mark.parametrize( + "mode", + [ + "missing_packet_rate", + "string_packet_rate", + "invalid_base64", + "orphan_audio", + "odd_packet", + ], +) +async def test_load_rejects_invalid_individual_audio_packets(tmp_path, sessions, mode): + # The concatenated PCM can still meet coverage/voicing thresholds, so the + # validator must also check each packet before accepting the whole run. + args = _args(tmp_path, "--sessions", str(sessions)) + async with _server(mode) as (url, states): + args.url = url + result = await asyncio.wait_for(driver.run(args), timeout=4) + assert result["ok"] is False + assert result["passed_sessions"] == 0 + rows = result["sessions"] + assert isinstance(rows, list) and len(rows) == sessions + assert all(row["error"] is not None for row in rows) + expected_code = { + "missing_packet_rate": "invalid_audio_sample_rate", + "string_packet_rate": "invalid_audio_sample_rate", + "invalid_base64": "invalid_audio_base64", + "orphan_audio": "missing_audio_response_id", + "odd_packet": "unaligned_pcm16_packet", + }[mode] + assert all({packet["code"] for packet in row["invalid_audio_packets"]} == {expected_code} for row in rows) + assert all( + sum(packet["samples"] for packet in row["audio_packet_timeline"]) == row["output_samples"] for row in rows + ) + assert all(state["closed"] and state["close_requested"] for state in states) + assert json.loads((tmp_path / "out/load-result.json").read_text()) == result + + +def test_load_allows_empty_and_multi_frame_audio_packets(): + client = driver.RawRealtimeProbe("ws://unused") + client.events.add(_audio(samples=0), received_at_s=1.0) + client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5), received_at_s=1.4) + metrics = driver._load_metrics(client, [(1.0, 1.0, 1.001)]) + assert metrics["invalid_audio_packets"] == [] + assert metrics["output_samples"] == driver.FRAME_SAMPLES * 5 + packets = metrics["audio_packet_timeline"] + assert isinstance(packets, list) and len(packets) == 1 + + +def test_load_packet_diagnostics_do_not_copy_remote_payload(): + client = driver.RawRealtimeProbe("ws://unused") + event = _audio() + event["delta"] = "private-remote-payload!" + client.events.add(event, received_at_s=1.4) + metrics = driver._load_metrics(client, []) + assert metrics["output_samples"] == 0 + assert metrics["invalid_audio_packets"] == [ + { + "event_index": 0, + "received_at_s": 1.4, + "code": "invalid_audio_base64", + } + ] + assert "private-remote-payload" not in json.dumps(metrics) From 8c3d3407d1bc9b1b538d2196f4ed2d802a15cfc3 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Thu, 17 Sep 2026 01:01:01 -0600 Subject: [PATCH 24/99] [CI] Fix Structured Structured Diffusion Config Wan Test (#7707) Signed-off-by: Alex Brooks --- tests/e2e/offline_inference/test_wan22_t2v.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py index 07ac5c8a6..b3c7d6aab 100644 --- a/tests/e2e/offline_inference/test_wan22_t2v.py +++ b/tests/e2e/offline_inference/test_wan22_t2v.py @@ -6,6 +6,7 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OfflineOmniClient from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" @@ -21,17 +22,19 @@ def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None: """The default WAN resolver launches the typed diffusion config itself.""" omni = omni_runner.omni resolved_stages = omni.engine.stage_configs + stage_clients = omni.engine.stage_clients assert len(resolved_stages) == 1 + assert len(stage_clients) == 1 stage_config = resolved_stages[0] + od_config = stage_clients[0].od_config assert isinstance(stage_config, VllmOmniDiffusionStageConfig) + assert isinstance(od_config, OmniDiffusionConfig) assert stage_config.stage_id == 0 assert stage_config.model_stage == "dit" assert stage_config.final_output_type == "video" assert omni.engine.stage_vllm_configs == [None] - od_config = omni.get_diffusion_od_config() - assert od_config is not None assert od_config.stage_id == 0 assert od_config.model == MODEL assert od_config.model_class_name == "WanPipeline" From cb439f3e84394c0cdfc252eaff3303fb099cc67d Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:41:52 +0800 Subject: [PATCH 25/99] [Diffusion][Model] Add direct FastH3 8-Step V2 checkpoint loading (#7610) Signed-off-by: david6666666 <530634352@qq.com> --- recipes/MiniMaxAI/MiniMax-H3.md | 62 ++ .../minimax_h3/test_fasth3_checkpoint.py | 679 ++++++++++++++++++ .../minimax_h3/test_minimax_h3_contract.py | 1 + .../minimax_h3/test_minimax_h3_fasth3.py | 5 +- .../test_minimax_h3_quantization.py | 23 +- .../test_minimax_h3_quantization_quality.py | 6 +- .../diffusion/models/minimax_h3/fasth3.py | 3 + .../models/minimax_h3/fasth3_checkpoint.py | 131 ++++ .../minimax_h3/minimax_h3_transformer.py | 76 +- .../models/minimax_h3/pipeline_minimax_h3.py | 52 +- .../models/minimax_h3/checkpoint.py | 29 +- 11 files changed, 1053 insertions(+), 14 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index b48a1f472..934e882f7 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1106,6 +1106,68 @@ generations move the end-to-end figure toward the denoising one. Fusing the adapter does not measurably change startup: weight loading took 77.3 s with it against 85.8 s without. +### FastH3 8-Step V2 full checkpoint + +[FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) +ships a full Diffusers-format transformer, including learned VSA compression +gates. Serve the Hugging Face model ID or a downloaded snapshot directly: + +```bash +export FASTH3_V2_MODEL=FastVideo/FastVideo-FastH3-8-Step-V2 +``` + +Use a Hugging Face account with access to the release, and pin `--revision` +when reproducing a result. The existing component loader reads its safetensors +shards; the H3 weight loader maps names and loads Q/K/V and MLP projections into +native parameters in memory. No conversion command or second transformer +checkpoint is needed. + +The transformer and text components come from the V2 release. To retain Omni's +native tiled and parallel VAE execution, the loader fetches only the video/audio +VAEs from the `MiniMaxAI/MiniMax-H3` revision pinned in V2's `provenance.json`. +Those components use the normal Hugging Face cache. No separate base-model +argument or full base-transformer download is required. Download time is a +preparation cost, separate from warmup and request latency. + +Install the optional `fastvideo-kernel` build required by the +`FASTVIDEO_VSA` backend, including its H3 block-map entry point. A four-GPU +configuration is: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "${FASTH3_V2_MODEL}" --omni \ + --host 127.0.0.1 --port 8095 --trust-remote-code --task-type fl2va \ + --served-model-name FastH3-V2 \ + --num-gpus 4 --usp 4 --ring 1 \ + --vae-patch-parallel-size 4 --vae-parallel-mode tile --vae-use-tiling \ + --diffusion-attention-backend FASTVIDEO_VSA + +curl --fail-with-body -sS http://127.0.0.1:8095/v1/videos/sync \ + -F 'model=FastH3-V2' \ + -F 'prompt=A red fox runs through fresh snow at dawn, with fast pawsteps, winter wind and distant birds.' \ + -F 'width=1344' -F 'height=768' -F 'aspect_ratio=16:9' -F 'fps=24' \ + -F 'num_inference_steps=8' -F 'guidance_scale=1' -F 'flow_shift=10' \ + -F 'seed=1101' \ + -F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \ + -o fasth3_v2.mp4 +``` + +V2 pins video/audio shifts **10/3**, guidance **1**, and the unshifted ladder +`[0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0]`. +Use **`num_inference_steps=8`** in Omni: these nine nodes bound eight model +evaluations. FastVideo's `--steps 9` counts nodes instead. The existing H3 ODE +update is reused without stochastic re-noising. + +The trained attention policy is **80% sparsity with 64-token video tiles**. +Omni derives the number of selected video tiles from each request's geometry; +do not supply a fixed `--fastvideo-vsa-topk`. Text/condition/audio prefix keys +remain available to every query, and prefix queries remain dense. Missing VSA +geometry or a failed H3 kernel raises an error instead of serving the student +with dense attention. + +This release supports T2VA only, with local or pure Ulysses attention. Additional +LoRA adapters, Ref2VA and FL2VA conditioning are unsupported. The legacy +four-step adapter keeps its original sampling and fixed-top-k behavior. + ## Key parameters | Parameter | Recommended value | Notes | diff --git a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py new file mode 100644 index 000000000..9a530bf90 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterable + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import ( + FASTH3_BASE_MODEL, + FASTH3_V2_BASE_SCHEDULE, + FASTH3_V2_MODEL_ID, + FastH3CheckpointSpec, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTModel, +) +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + resolve_minimax_h3_diffusion_model_path, +) +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest +from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + is_minimax_h3_modular, + resolve_minimax_h3_model_root, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _fastvideo_metadata() -> dict[str, object]: + """Return the release-owned metadata read from fastvideo_inference.json.""" + return { + "model_id": FASTH3_V2_MODEL_ID, + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.8, + "vsa_tile_size": 64, + "task": "t2av", + } + + +def _sampling( + *, + num_inference_steps: int | str | None = 8, + guidance_scale: float | None = 1.0, + extra_args: dict[str, object] | None = None, + lora_request: LoRARequest | None = None, + timesteps: torch.Tensor | None = None, + sigmas: list[float] | None = None, +) -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + # Keep the string form to exercise malformed client input validation. + num_inference_steps=num_inference_steps, # type: ignore[arg-type] + guidance_scale=guidance_scale, + extra_args=extra_args or {}, + lora_request=lora_request, + timesteps=timesteps, + sigmas=sigmas, + ) + + +def _od_config( + *, + backend: str = "FASTVIDEO_VSA", + topk: int | None = None, + per_role: dict[str, AttentionSpec] | None = None, + lora_path: str | None = None, + ring_degree: int = 1, + allgather_degree: int = 1, +) -> OmniDiffusionConfig: + attention_config = AttentionConfig( + default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk), + per_role=per_role or {}, + ) + return OmniDiffusionConfig( + diffusion_attention_config=attention_config, + lora_path=lora_path, + parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree), + ) + + +def test_v2_release_has_nine_sigma_nodes_and_eight_intervals(): + assert FASTH3_V2_BASE_SCHEDULE.base_schedule == ( + 0.999, + 0.874, + 0.749, + 0.624, + 0.5, + 0.375, + 0.25, + 0.125, + 0.0, + ) + assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9 + assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8 + assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx( + [ + 0.9998999099189271, + 0.985788405143244, + 0.9675752486758817, + 0.94316807738815, + 0.9090909090909091, + 0.8571428571428571, + 0.7692307692307693, + 0.5882352941176471, + 0.0, + ], + abs=1e-12, + ) + assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx( + [ + 0.9996664442961973, + 0.9541484716157204, + 0.8995196156925539, + 0.8327402135231315, + 0.75, + 0.6428571428571429, + 0.5, + 0.3, + 0.0, + ], + abs=1e-12, + ) + + +def test_fastvideo_metadata_builds_release_sampling_metadata(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + + assert spec.vsa_sparsity == 0.8 + assert spec.release_metadata() == { + "partition": "fl2va", + "tasks": ["t2va"], + "sigma_shift_scales": {"video": 10.0, "audio": 3.0}, + "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule), + } + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("model_id", "FastVideo/FastH3-other"), + ("schema_version", "fasth3-inference-contract-v0"), + ("dmd_denoising_steps", [999, 749, 500, 250]), + ("transformer_forwards", 4), + ("video_scheduler_shift", 12.0), + ("audio_scheduler_shift", 4.0), + ("guidance_scale", 2.0), + ("attention_backend", "FLASH_ATTN"), + ("vsa_sparsity", 0.9), + ("vsa_tile_size", 32), + ("task", "t2va"), + ], +) +def test_fastvideo_metadata_rejects_contract_drift(field, value): + metadata = _fastvideo_metadata() + metadata[field] = value + + with pytest.raises(ValueError, match=field): + FastH3CheckpointSpec.from_metadata(metadata) + + +def test_fastvideo_metadata_rejects_missing_contract_field(): + metadata = _fastvideo_metadata() + metadata.pop("task") + + with pytest.raises(ValueError, match="task"): + FastH3CheckpointSpec.from_metadata(metadata) + + +def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release(): + with pytest.raises(ValueError, match="model_id"): + FastH3CheckpointSpec.from_metadata( + { + "schema_version": 1, + "partition": "fl2va", + "tasks": ["t2va", "fl2va"], + "sigma_shift_scales": {"video": 12.0, "audio": 3.0}, + } + ) + + +def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role(): + config = _od_config( + backend="TORCH_SDPA", + per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")}, + ) + FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract( + partition="fl2va", od_config=config + ) + + +@pytest.mark.parametrize( + ("kwargs", "partition", "match"), + [ + ({}, "ref2va", "task-type fl2va"), + ({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"), + ({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"), + ({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"), + ({"ring_degree": 2}, "fl2va", "local attention"), + ({"allgather_degree": 2}, "fl2va", "local attention"), + ], +) +def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(ValueError, match=match): + spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs)) + + +def test_v2_serving_contract_reads_a_per_role_fixed_topk_override(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + config = _od_config( + per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)}, + ) + with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"): + spec.check_serving_contract(partition="fl2va", od_config=config) + + +def test_v2_request_contract_accepts_omitted_steps_and_default_shifts(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + spec.check_request(_sampling(num_inference_steps=None)) + spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1)) + + +@pytest.mark.parametrize( + "sampling", + [ + _sampling(num_inference_steps=None), + _sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])), + _sampling(sigmas=[1.0, 0.5, 0.0]), + ], +) +def test_v2_step_execution_requires_pinned_count_and_schedule(sampling): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(OmniClientError, match="step execution"): + spec.check_request(sampling, step_execution=True) + + +def test_v2_step_execution_accepts_only_explicit_pinned_count(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + spec.check_request(_sampling(num_inference_steps=8), step_execution=True) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_inference_steps": 7}, "num_inference_steps=8"), + ({"num_inference_steps": 9}, "num_inference_steps=8"), + ({"num_inference_steps": "8"}, "num_inference_steps=8"), + ({"extra_args": {"flow_shift": 9}}, "flow_shift=10"), + ({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"), + ({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"), + ({"guidance_scale": 2.0}, "guidance_scale=1"), + ( + {"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")}, + "per-request LoRA", + ), + ], +) +def test_v2_request_contract_rejects_sampling_drift(kwargs, match): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(OmniClientError, match=match): + spec.check_request(_sampling(**kwargs)) + + +def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch): + model_root = tmp_path / "FastH3" + model_root.mkdir() + revision = "base-revision-123" + (model_root / "provenance.json").write_text( + json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}), + encoding="utf-8", + ) + snapshot = tmp_path / "hf-snapshot" + calls: dict[str, object] = {} + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr( + "vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific", + fake_download, + ) + + result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root) + + assert result == snapshot / "FL2VA" + assert calls == { + "model_name_or_path": FASTH3_BASE_MODEL, + "cache_dir": None, + "allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"], + "revision": revision, + "require_all": True, + } + assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"] + assert not (model_root / "model_index.json").exists() + assert not (model_root / "conversion.json").exists() + + +@pytest.mark.parametrize( + "base_model", + [ + "MiniMaxAI/MiniMax-H3@main", + "hf://MiniMaxAI/MiniMax-H3@", + "hf://Other/MiniMax-H3@revision", + "", + None, + ], +) +def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model): + model_root = tmp_path / "FastH3" + model_root.mkdir() + (model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8") + + with pytest.raises(ValueError, match="provenance"): + FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root) + + +def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path): + root = tmp_path / "FastH3" + (root / "text_encoder").mkdir(parents=True) + (root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8") + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert is_minimax_h3_modular(str(root)) is True + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder") + assert not (root / "FL2VA").exists() + + +def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path): + from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module + + snapshot = tmp_path / "modular-snapshot" + calls: dict[str, object] = {} + + def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]: + assert model == "FastVideo/FastVideo-FastH3-8-Step-V2" + assert revision == "release-revision" + return {"_class_name": "MiniMaxH3ModularPipeline"} + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index) + monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download) + + result = resolve_minimax_h3_model_root( + "FastVideo/FastVideo-FastH3-8-Step-V2", + "release-revision", + "fl2va", + ) + + assert result == str(snapshot / "text_encoder") + assert calls == { + "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2", + "cache_dir": None, + "allow_patterns": ["text_encoder/**"], + "revision": "release-revision", + "require_all": True, + } + + +def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch): + from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module + + monkeypatch.setattr( + checkpoint_module, + "get_diffusion_model_index", + lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"}, + ) + + assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False + + +def test_native_local_model_root_keeps_partition_layout(tmp_path): + root = tmp_path / "MiniMax-H3" + (root / "FL2VA" / "text_encoder").mkdir(parents=True) + (root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8") + + assert is_minimax_h3_modular(str(root)) is False + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") + + +def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path): + root = tmp_path / "FastH3" + root.mkdir() + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root) + assert not (root / "FL2VA").exists() + assert not (root / "Ref2VA").exists() + + +def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + snapshot = tmp_path / "modular-snapshot" + calls: dict[str, object] = {} + + def fake_modular(_model: str, _revision: str | None = None) -> bool: + return True + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular) + monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download) + + result = resolve_minimax_h3_diffusion_model_path( + "FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va" + ) + + assert result == str(snapshot) + assert calls == { + "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2", + "cache_dir": None, + "allow_patterns": [ + "modular_model_index.json", + "fastvideo_inference.json", + "provenance.json", + "transformer/**", + ], + "revision": "release-revision", + "require_all": True, + } + patterns = calls["allow_patterns"] + assert isinstance(patterns, list) + assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns) + + +def test_diffusers_arch_config_normalizes_all_aliases(): + arch = MiniMaxH3DiTArchConfig.from_mapping( + { + "num_refiner_layers": 3, + "ffn_dim": 12, + "in_channels": 4, + "audio_in_channels": 6, + "freq_dim": 8, + "time_embed_hidden_dim": 10, + "rope_freq_dim": 5, + "patch_size": [1, 2, 2], + } + ) + + assert arch.token_refiner_num_layers == 3 + assert arch.ffn_hidden_size == 12 + assert arch.latents_dim == 4 + assert arch.audio_latents_dim == 6 + assert arch.timestep_input_dim == 8 + assert arch.time_embed_hidden_size == 10 + assert arch.rope_inv_freq_len == 5 + assert arch.patch_size == (1, 2, 2) + + +class _WeightTarget(nn.Module): + def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(shape)) + self.weight.weight_loader = loader + + +class _PipelineLoadTransformer(nn.Module): + def __init__(self, loaded: set[str]) -> None: + super().__init__() + self.arch = MiniMaxH3DiTArchConfig(num_layers=1) + self._loaded = loaded + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + list(weights) + return self._loaded + + def post_load_weights(self) -> None: + pass + + +def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"}) + pipeline.video_vae = None + pipeline.audio_vae = None + pipeline._fasth3 = None + pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + + with pytest.raises(ValueError, match="missing compression gates"): + pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))]) + + +def _load_fixture( + *, diffusers_weights: bool +) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]: + model = object.__new__(MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.arch = MiniMaxH3DiTArchConfig( + num_layers=1, + hidden_size=3, + num_attention_heads=2, + attention_head_dim=2, + ffn_hidden_size=6, + rope_inv_freq_len=2, + ) + model._diffusers_weights = diffusers_weights + model._rope_theta = 10000.0 + qkv_calls: list[tuple[str | None, torch.Tensor]] = [] + fc1_calls: list[tuple[int, torch.Tensor]] = [] + + def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None: + if shard_id is None: + assert loaded_weight.shape == (12, 3) + param.data.copy_(loaded_weight) + qkv_calls.append((None, loaded_weight.clone())) + return + assert loaded_weight.shape == (4, 3) + assert shard_id in {"q", "k", "v"} + offset = {"q": 0, "k": 4, "v": 8}[shard_id] + param.data[offset : offset + 4].copy_(loaded_weight) + qkv_calls.append((shard_id, loaded_weight.clone())) + + def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None: + assert loaded_weight.shape == (3, 3) + param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight) + fc1_calls.append((shard_id, loaded_weight.clone())) + + model.blocks = nn.ModuleList([nn.Module()]) + model.blocks[0].attn = nn.Module() + model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader) + model.blocks[0].mlp = nn.Module() + model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader) + model.rope = nn.Module() + model.rope.register_buffer("inv_freq", torch.zeros(2)) + return model, qkv_calls, fc1_calls + + +def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices(): + model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True) + q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10 + k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100 + v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200 + mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300 + + loaded = model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", q), + ("transformer_blocks.0.attn.to_k.weight", k), + ("transformer_blocks.0.attn.to_v.weight", v), + ("transformer_blocks.0.ff.net.0.proj.weight", mlp), + ] + ) + + assert loaded == { + "blocks.0.attn.qkv_proj.weight", + "blocks.0.mlp.fc1.weight", + "rope.inv_freq", + } + assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"] + assert [shard_id for shard_id, _ in fc1_calls] == [0, 1] + torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v))) + torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3]))) + torch.testing.assert_close( + model.rope.inv_freq, + 1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)), + rtol=0, + atol=0, + ) + + +def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts(): + model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False) + grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3) + fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3) + + loaded = model.load_weights( + [ + ("blocks.0.attn.qkv_proj.weight", grouped_qkv), + ("blocks.0.mlp.fc1.weight", fc1), + ] + ) + + expected_qkv = torch.cat( + ( + grouped_qkv[:2], + grouped_qkv[6:8], + grouped_qkv[2:4], + grouped_qkv[8:10], + grouped_qkv[4:6], + grouped_qkv[10:], + ) + ) + assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"} + assert [shard_id for shard_id, _ in qkv_calls] == [None] + assert [shard_id for shard_id, _ in fc1_calls] == [0, 1] + torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv) + torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1) + assert torch.equal(model.rope.inv_freq, torch.zeros(2)) + + +def test_diffusers_load_weights_rejects_unknown_name(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"): + model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))]) + + +def test_diffusers_load_weights_rejects_duplicate_source_name(): + model, _, _ = _load_fixture(diffusers_weights=True) + weight = torch.zeros(4, 3) + + with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"): + model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", weight), + ("transformer_blocks.0.attn.to_q.weight", weight), + ] + ) + + +def test_diffusers_load_weights_rejects_incomplete_qkv_group(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"): + model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)), + ("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)), + ] + ) + + +def test_diffusers_load_weights_rejects_missing_compression_gate_parameter(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="has no model parameter"): + model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))]) + + +def test_resolve_sigma_positions_uses_release_metadata_schedule(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + nn.Module.__init__(pipeline) + pipeline._fasth3 = None + pipeline._lora_sigma_schedules = {} + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + pipeline._base_schedule_by_partition = { + "fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()), + } + + positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling()) + assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule + assert steps == 8 == len(positions) - 1 + with pytest.raises(OmniClientError, match="must be 8"): + pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9)) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 4b6fc0d88..cea4862d1 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( assert pipeline.partition == expected_partition assert pipeline.supported_tasks == expected_tasks assert len(created["dit"]) == expected_dits + assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"]) component_path = tmp_path / component_partition assert created["video_vae"] == [str(component_path / "video_vae")] assert created["audio_vae"] == [str(component_path / "audio_vae")] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py index 5c874ff02..c89c37db9 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py @@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path): def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path): path = tmp_path / "fasth3" / "adapter_model.safetensors" - _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}) + _write_adapter( + path, + tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, + ) # Dropping it silently would load a model that is not the distilled student. with pytest.raises(FastH3AdapterError, match="name no known"): _load(path.parent) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index 92bd09ac2..e1f83fd01 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace from unittest.mock import Mock @@ -41,7 +42,7 @@ class _FakeAttention(nn.Module): super().__init__() -def _small_od_config(): +def _small_od_config(*, class_name=None): arch = { "num_layers": 1, "token_refiner_num_layers": 1, @@ -60,6 +61,8 @@ def _small_od_config(): "final_adaln_out_features": 2 * 8, "rope_inv_freq_len": 2, } + if class_name is not None: + arch["_class_name"] = class_name return SimpleNamespace( tf_model_config=arch, parallel_config=SimpleNamespace(ulysses_degree=1), @@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch): ) +def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1) + + model = h3.MiniMaxH3DiTModel( + _small_od_config(class_name="MiniMaxH3Transformer3DModel"), + diffusers_weights=False, + ) + + assert model._diffusers_weights is False + + class _WeightTarget(nn.Module): def __init__(self, loader): super().__init__() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 87b04613a..3beda5acc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -187,7 +187,11 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig): "model": model_ref, "enforce_eager": True, # The fused BF16 baseline only fits on H100-80GB with encoder TP. - "stage_0_tensor_parallel_size": 2, + # This test loads the direct FL2VA diffusion pipeline (not the + # disaggregated two-stage MiniMax-H3 pipeline), so configure both the + # DiT world and the Qwen3-VL encoder TP explicitly. + "tensor_parallel_size": 2, + "text_encoder_tp_size": 2, "vae_use_tiling": True, # MiniMax H3's VAE is implemented by code shipped with the checkpoint. "trust_remote_code": True, diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py index c4860c57e..dcc513028 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -102,6 +102,7 @@ _MODEL_LEVEL_TARGETS = { "time_embedder.linear_2": "time_embedder.proj_out", "norm_out.linear": "final_layer.adaln_proj.linear", "norm_out.norm": "final_layer.norm", + "token_refiner.final_norm": "token_refiner.final_norm", } # Per-block adapter suffix -> (native suffix, how a delta enters the native @@ -115,6 +116,8 @@ _BLOCK_TARGETS = { "attn.to_v": ("attn.qkv_proj", _QKV_V), "attn.to_out.0": ("attn.out_proj", _PLAIN), "attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN), + "attn.norm_q": ("attn.q_norm", _PLAIN), + "attn.norm_k": ("attn.k_norm", _PLAIN), "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES), "ff.net.2": ("mlp.fc2", _PLAIN), "adaln_proj.linear": ("adaln_proj.linear", _PLAIN), diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py new file mode 100644 index 000000000..0382e9220 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Sampling contract for FastVideo FastH3 full checkpoints.""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2" +FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0)) + + +@dataclass(frozen=True) +class FastH3CheckpointSpec: + """The full V2 release has its own schedule and trained attention policy. + + This is independent of ``FastH3WeightFusion``: the released weights already + contain the entire student, including its learned compression gates. + """ + + vsa_sparsity: float = 0.8 + + @classmethod + def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec: + """Validate the release's own fastvideo_inference.json.""" + required = { + "model_id": FASTH3_V2_MODEL_ID, + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.8, + "vsa_tile_size": 64, + "task": "t2av", + } + for key, expected in required.items(): + if metadata.get(key) != expected: + raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}") + return cls() + + def release_metadata(self) -> dict[str, Any]: + """Express the sampling policy in the existing H3 pipeline schema.""" + return { + "partition": "fl2va", + "tasks": ["t2va"], + "sigma_shift_scales": {"video": 10.0, "audio": 3.0}, + "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule), + } + + def resolve_native_vaes(self, model_root: Path) -> Path: + """Reuse the frozen base VAEs with Omni's native tiled/parallel runtime. + + Only VAE components are fetched; the student and text encoder come + directly from the FastVideo release. The release pins the base revision. + """ + provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8")) + base = provenance.get("base_model", "") + prefix = f"hf://{FASTH3_BASE_MODEL}@" + if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]: + raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision") + return ( + Path( + download_weights_from_hf_specific( + model_name_or_path=FASTH3_BASE_MODEL, + cache_dir=None, + allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"], + revision=base[len(prefix) :], + require_all=True, + ) + ) + / "FL2VA" + ) + + def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None: + if partition != "fl2va": + raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)") + if getattr(od_config, "lora_path", None): + raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported") + if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA": + raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA") + attention_config = getattr(od_config, "diffusion_attention_config", None) + per_role = getattr(attention_config, "per_role", None) or {} + spec = per_role.get("self") or getattr(attention_config, "default", None) + if getattr(spec, "fastvideo_vsa_topk", None) is not None: + raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override") + parallel = getattr(od_config, "parallel_config", None) + if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")): + raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism") + + def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None: + if step_execution: + # StepScheduler admits requests before the pipeline hook runs and + # derives its lifetime from these fields. The pinned schedule must + # therefore be explicit and cannot be overridden by custom arrays. + if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps: + raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8") + if sampling.timesteps is not None or sampling.sigmas is not None: + raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas") + if sampling.lora_request is not None: + raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters") + steps = sampling.num_inference_steps + if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps: + raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted") + extra = sampling.extra_args or {} + for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)): + try: + value = float(extra.get(key, expected)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc + if not math.isclose(value, expected): + raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}") + if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0: + raise OmniClientError("FastH3 V2 requires guidance_scale=1") diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index e26493931..5662a9641 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -55,6 +55,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.platforms import current_omni_platform +from .fasth3 import _resolve_native_target + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -118,6 +120,18 @@ class MiniMaxH3DiTArchConfig: @classmethod def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig: + # The modular Diffusers checkpoint uses these spellings for the same + # architecture. Normalize before constructing the existing native DiT. + aliases = { + "num_refiner_layers": "token_refiner_num_layers", + "ffn_dim": "ffn_hidden_size", + "in_channels": "latents_dim", + "audio_in_channels": "audio_latents_dim", + "freq_dim": "timestep_input_dim", + "time_embed_hidden_dim": "time_embed_hidden_size", + "rope_freq_dim": "rope_inv_freq_len", + } + config = {aliases.get(name, name): value for name, value in config.items()} fields = cls.__dataclass_fields__ values = {name: config[name] for name in fields if name in config} if "patch_size" in values: @@ -394,6 +408,9 @@ def _sdpa_varlen_attention( class MiniMaxH3Attention(nn.Module): + # Full sparse checkpoints pin a ratio; legacy adapters use backend top-k. + vsa_sparsity: float | None = None + def __init__( self, arch: MiniMaxH3DiTArchConfig, @@ -588,6 +605,7 @@ class MiniMaxH3Attention(nn.Module): # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE, + **({"vsa_h3_sparsity": self.vsa_sparsity} if self.vsa_sparsity is not None else {}), # Present only for a VSA artifact; the VSA backend reads it as # the learned compression gate and every other backend ignores it. **({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}), @@ -1127,10 +1145,23 @@ class MiniMaxH3DiTModel(nn.Module): self, od_config: OmniDiffusionConfig, quant_config: QuantizationConfig | None = None, + *, + diffusers_weights: bool | None = None, ) -> None: super().__init__() tf_config = od_config.tf_model_config config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config) + # The native MiniMax-H3 Hub snapshot advertises the Diffusers + # transformer class in its root config, while its FL2VA/Ref2VA + # components still contain native weights. The pipeline has already + # resolved the actual source format, so let it override the + # class-name heuristic. Keep the heuristic for standalone callers. + self._diffusers_weights = ( + config_mapping.get("_class_name") == "MiniMaxH3Transformer3DModel" + if diffusers_weights is None + else diffusers_weights + ) + self._rope_theta = float(config_mapping.get("rope_theta", 10000.0)) arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping) self.arch = arch self.od_config = od_config @@ -1211,7 +1242,7 @@ class MiniMaxH3DiTModel(nn.Module): ) self._mark_missing_params_required() - def enable_vsa_gates(self) -> None: + def enable_vsa_gates(self, *, sparsity: float | None = None) -> None: """Give every DiT block's attention a VSA compression gate. A FastH3 VSA artifact assigns these projections rather than adding to @@ -1223,6 +1254,9 @@ class MiniMaxH3DiTModel(nn.Module): return for block in self.blocks: block.attn.enable_vsa_gate() + block.attn.vsa_sparsity = sparsity + if sparsity is not None: + block.attn.to_gate_compress.weight.missing_param_init = "error" self.vsa_gates_enabled = True def _mark_missing_params_required(self) -> None: @@ -1293,17 +1327,37 @@ class MiniMaxH3DiTModel(nn.Module): self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - """Load exact H3 checkpoint names with logical TP-aware loaders.""" + """Load native or Diffusers H3 weights with the existing TP-aware loaders.""" params = dict(self.named_parameters()) params.update(dict(self.named_buffers())) loaded: set[str] = set() + diffusers_weights = getattr(self, "_diffusers_weights", False) + qkv_parts: dict[str, set[str]] = {} + source_names: set[str] = set() for name, loaded_weight in weights: + layout = "plain" + if diffusers_weights: + if name in source_names: + raise ValueError(f"duplicate Diffusers H3 weight: {name}") + source_names.add(name) + module, _, kind = name.rpartition(".") + target = _resolve_native_target(module) + if target is None or kind not in {"weight", "bias"}: + raise ValueError(f"unsupported Diffusers H3 weight: {name}") + name, layout = f"{target[0]}.{kind}", target[1] param = params.get(name) if param is None: + if diffusers_weights: + raise ValueError(f"Diffusers H3 weight has no model parameter: {name}") logger.warning("Skipping MiniMax H3 weight not present in model: %s", name) continue weight_loader = getattr(param, "weight_loader", default_weight_loader) - if name.endswith(".attn.qkv_proj.weight"): + if layout in {"q", "k", "v"}: + # vLLM can load each projection directly into its packed QKV + # parameter, including TP slicing and online quantization. + weight_loader(param, loaded_weight, layout) + qkv_parts.setdefault(name, set()).add(layout) + elif name.endswith(".attn.qkv_proj.weight"): # Transform checkpoint layout before entering vLLM's loader so # online FP8 can keep ``online_process_loader`` outermost. loaded_weight = _reorder_grouped_qkv_to_qkv( @@ -1319,12 +1373,26 @@ class MiniMaxH3DiTModel(nn.Module): "MiniMax H3 fc1 checkpoint rows must split evenly into " f"gate/up matrices, got {tuple(loaded_weight.shape)}" ) - gate, up = loaded_weight.chunk(2, dim=0) + first, second = loaded_weight.chunk(2, dim=0) + gate, up = (second, first) if layout == "swap_halves" else (first, second) weight_loader(param, gate, 0) weight_loader(param, up, 1) else: weight_loader(param, loaded_weight) loaded.add(name) + if diffusers_weights: + for name, parts in qkv_parts.items(): + if parts != {"q", "k", "v"}: + raise ValueError(f"incomplete Diffusers H3 QKV group {name}: {sorted(parts)}") + # Diffusers reconstructs RoPE from config instead of storing this + # native checkpoint buffer. Compute on CPU for identical values. + freq_dim = self.arch.rope_inv_freq_len + rope = 1.0 / ( + self._rope_theta + ** (torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device="cpu") / (2 * freq_dim)) + ) + default_weight_loader(params["rope.inv_freq"], rope) + loaded.add("rope.inv_freq") return loaded @staticmethod diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 1b8bd7fea..60d6f4125 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -67,6 +67,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + is_minimax_h3_modular, resolve_minimax_h3_partition, ) from vllm_omni.model_executor.models.minimax_h3.conditioning import ( @@ -106,6 +107,7 @@ from .denoise_loop import ( ) from .encoder import MiniMaxH3Qwen3VLEncoder from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion +from .fasth3_checkpoint import FastH3CheckpointSpec from .lora import TurboSpec, load_minimax_h3_turbo_lora from .minimax_h3_transformer import ( MiniMaxH3Attention, @@ -213,7 +215,11 @@ def _resolve_minimax_h3_model_root( if path.name in {"FL2VA", "Ref2VA"}: return path.parent return path - if load_text_encoder: + if is_minimax_h3_modular(model, revision): + allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"] + if load_text_encoder: + allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"] + elif load_text_encoder: allow_patterns = ( MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition] ) @@ -314,6 +320,8 @@ def resolve_minimax_h3_diffusion_model_path( partition, load_text_encoder=False, ) + if is_minimax_h3_modular(str(model_root), revision): + return str(model_root) if partition == "combined": return str(model_root) subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" @@ -575,6 +583,7 @@ class MiniMaxH3Pipeline( _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} # Set from --lora-path during construction; absent means no FastH3 adapter. _fasth3: FastH3WeightFusion | None = None + _fasth3_checkpoint: FastH3CheckpointSpec | None = None def _load_diffusion_lora_adapter( self, @@ -800,10 +809,13 @@ class MiniMaxH3Pipeline( self._PROFILER_TARGETS.remove("encode_prompt") if not self.load_vae_encoder: self._PROFILER_TARGETS.remove("_encode_local_media") + modular = is_minimax_h3_modular(str(od_config.model), od_config.revision) self.partition = _minimax_h3_partition_for_task( getattr(od_config, "task_type", None), str(od_config.model), ) + if modular and str(od_config.task_type or "auto").lower() == "auto": + self.partition = "fl2va" self._turbo_lora_specs: dict[int, TurboSpec] = {} self._native_lora_adapter_ids: set[int] = set() self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {} @@ -813,9 +825,19 @@ class MiniMaxH3Pipeline( self.partition, load_text_encoder=self.load_text_encoder, ) - model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") - model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) - release = model_index.get("_minimax_h3") or {} + if modular: + model_path = model_root + self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata( + json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8")) + ) + self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config) + release = self._fasth3_checkpoint.release_metadata() + vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root) + else: + model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") + model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) + release = model_index.get("_minimax_h3") or {} + vae_model_path = model_path partition = str(release.get("partition", "")).lower() expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va" if partition != expected_partition: @@ -876,11 +898,18 @@ class MiniMaxH3Pipeline( self.transformer = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + diffusers_weights=modular, ) + if self._fasth3_checkpoint is not None: + self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity) + logger.info( + "FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64" + ) if ref2va_model_path is not None: self.transformers_ref = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + diffusers_weights=modular, ) self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer) @@ -963,14 +992,14 @@ class MiniMaxH3Pipeline( # so VAEs stay resident for new configurations. component_load_device = torch.device("cpu") if legacy_manual_components else self.device self.video_vae = MiniMaxH3VideoVAE( - os.path.join(model_path, "video_vae"), + os.path.join(vae_model_path, "video_vae"), device=self.device, load_device=component_load_device, decode_only=not self.load_vae_encoder, trust_remote_code=od_config.trust_remote_code, ) self.audio_vae = MiniMaxH3AudioVAE( - os.path.join(model_path, "audio_vae"), + os.path.join(vae_model_path, "audio_vae"), device=self.device, load_device=component_load_device, decode_only=not self.load_vae_encoder, @@ -1046,6 +1075,12 @@ class MiniMaxH3Pipeline( # load_weights only warns on a parameter the model does not have, so # close the adapter against what the DiT actually consumed. self._fasth3.validate_fully_applied(transformer_loaded) + if self._fasth3_checkpoint is not None: + required_gates = { + f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers) + } + if missing_gates := required_gates - transformer_loaded: + raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}") return loaded_with_prefix @property @@ -2056,6 +2091,11 @@ class MiniMaxH3Pipeline( self._validate_turbo_sampling(sampling, turbo_spec) if has_native_lora: self._validate_native_sampling(sampling, task=task) + if self._fasth3_checkpoint is not None: + self._fasth3_checkpoint.check_request( + sampling, + step_execution=bool(getattr(self.od_config, "step_execution", False)), + ) if self._fasth3 is not None: self._fasth3.check_request( sampling, diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py index 0143e9675..b39af5cf8 100644 --- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -1,10 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 checkpoint resolution for the AR text-encoder stage.""" + from pathlib import Path +from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +_FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/" + MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { partition: [ f"{subdir}/text_encoder/**", @@ -16,6 +22,22 @@ MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"] +def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool: + path = Path(model) + if path.is_dir(): + # Native MiniMax-H3 snapshots can advertise the Diffusers modular + # class in ``model_index.json`` while retaining the legacy FL2VA / + # Ref2VA layout. FastH3 modular releases carry this marker file. + return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file() + # The FastVideo release is the modular checkpoint handled by this + # resolver. Do not classify the native MiniMax-H3 Hub repository from + # its class name alone; its partition files are still required here. + if not model.startswith(_FASTH3_MODULAR_MODEL_PREFIX): + return False + index = get_diffusion_model_index(model, revision=revision) or {} + return index.get("_class_name") == "MiniMaxH3ModularPipeline" + + def resolve_minimax_h3_partition( model: str, task_type: str | None, @@ -42,10 +64,13 @@ def resolve_minimax_h3_encoder_model_root( ) -> str: path = Path(model) partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va") + modular = is_minimax_h3_modular(model, revision) if path.is_dir(): if path.name == "text_encoder" and (path / "config.json").is_file(): return str(path) + if modular: + return str(path / "text_encoder") if path.name in {"FL2VA", "Ref2VA"}: path = path.parent subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" @@ -53,11 +78,13 @@ def resolve_minimax_h3_encoder_model_root( snapshot = download_weights_from_hf_specific( model_name_or_path=model, cache_dir=None, - allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition], + allow_patterns=["text_encoder/**"] if modular else MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition], revision=revision, require_all=True, ) + if modular: + return str(Path(snapshot) / "text_encoder") subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" return str(Path(snapshot) / subdir / "text_encoder") From 825ab20b07268050015516f9a65c733d7d089d28 Mon Sep 17 00:00:00 2001 From: MrDongsls Date: Thu, 17 Sep 2026 16:50:40 +0800 Subject: [PATCH 26/99] [Model][PersonaPlex] add per-row RingKV offsets and active rows (#7670) Signed-off-by: Mr.Dong <1042542469@qq.com> Co-authored-by: Mr.Dong <1042542469@qq.com> --- .../personaplex/test_ring_per_row_offsets.py | 369 ++++++++++++++++++ .../models/personaplex/personaplex_mimi.py | 94 +++-- .../personaplex/personaplex_temporal.py | 76 ++-- 3 files changed, 483 insertions(+), 56 deletions(-) create mode 100644 tests/model_executor/models/personaplex/test_ring_per_row_offsets.py diff --git a/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py new file mode 100644 index 000000000..fce42f40c --- /dev/null +++ b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from pytest_mock import MockerFixture + +from vllm_omni.model_executor.models.personaplex.personaplex_mimi import ( + PersonaPlexMimiCodec, + _MimiStreamingTransformer, + _StreamConv1d, + _StreamConvTr1d, +) +from vllm_omni.model_executor.models.personaplex.personaplex_temporal import ( + PersonaPlexTemporalStreaming, + _RingKV, +) + +pytestmark = pytest.mark.core_model + +SEED = 1234 +CUDA_DEVICE = torch.device("cuda") + + +def _mask(*rows: bool, device: torch.device | str = "cpu") -> torch.Tensor: + return torch.tensor(rows, dtype=torch.bool, device=device) + + +def _inputs(shape: tuple[int, ...], seed: int, device: torch.device) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(*shape, generator=generator).to(device) + + +def _make_mimi_transformer( + batch_size: int, + device: torch.device, + context: int = 6, +) -> _MimiStreamingTransformer: + torch.manual_seed(SEED) + transformer = _MimiStreamingTransformer( + num_layers=2, + dim=32, + num_heads=4, + context=context, + ).to(device) + for parameter in transformer.parameters(): + if parameter.dtype.is_floating_point: + nn.init.normal_(parameter, std=0.1) + transformer.streaming_init(batch_size) + return transformer + + +def _make_temporal( + batch_size: int, + device: torch.device, + context: int = 6, +) -> PersonaPlexTemporalStreaming: + torch.manual_seed(SEED) + temporal = PersonaPlexTemporalStreaming( + dim=16, + num_layers=2, + num_heads=4, + hidden=32, + context=context, + text_card=11, + ).to(device) + for parameter in temporal.parameters(): + if parameter.dtype.is_floating_point: + nn.init.normal_(parameter, std=0.1) + temporal.streaming_init(batch_size) + return temporal + + +def _assert_valid_ring_row_matches( + batched: _RingKV, + reference: _RingKV, + batched_positions: torch.Tensor, + reference_positions: torch.Tensor, + row: int, +) -> None: + assert torch.equal(batched_positions[row], reference_positions) + assert torch.equal(batched.end_offset[row], reference.end_offset[0]) + + # The active mask covers the physical write as well as offset advancement, + # so the complete per-row ring state remains singleton-identical. + assert torch.equal(batched.cache[:, row], reference.cache[:, 0]) + + +def _stream_carry(stream: _StreamConv1d | _StreamConvTr1d) -> torch.Tensor: + if isinstance(stream, _StreamConv1d): + assert stream.prev is not None + return stream.prev + assert stream.partial is not None + return stream.partial + + +def _make_passthrough_stage(mocker: MockerFixture): + stage = mocker.Mock(spec=["reset", "reset_slot"]) + stage.side_effect = lambda x, active: x + return stage + + +def _make_passthrough_transformer(mocker: MockerFixture): + transformer = mocker.Mock(spec=["streaming_init", "reset_streaming", "reset_slot", "step"]) + transformer.step.side_effect = lambda x, active: x + return transformer + + +def _active_args(call_args_list) -> torch.Tensor: + return torch.stack([call.args[1] for call in call_args_list]) + + +def _make_codec_stub(mocker: MockerFixture) -> PersonaPlexMimiCodec: + codec = PersonaPlexMimiCodec.__new__(PersonaPlexMimiCodec) + nn.Module.__init__(codec) + + quantizer = mocker.Mock(spec=["encode", "decode"]) + quantizer.encode.side_effect = lambda x: torch.zeros(8, x.shape[0], x.shape[-1], dtype=torch.long) + quantizer.decode.side_effect = lambda codes: torch.zeros(codes.shape[0], 1, codes.shape[-1]) + + codec.device = torch.device("cpu") + codec.dtype = torch.float32 + codec.model = SimpleNamespace(quantizer=quantizer) + codec._enc_stages = [] + codec._dec_stages = [] + codec._downsample = _make_passthrough_stage(mocker) + codec._upsample = _make_passthrough_stage(mocker) + codec.encoder_transformer = _make_passthrough_transformer(mocker) + codec.decoder_transformer = _make_passthrough_transformer(mocker) + codec.streaming_init(batch_size=2) + return codec + + +@pytest.mark.cpu +def test_codec_entrypoints_forward_active_mask(mocker: MockerFixture) -> None: + codec = _make_codec_stub(mocker) + active = _mask(True, False) + all_active = torch.ones_like(active) + + codec.encode_frame(torch.zeros(2, 1920), active) + codec.decode_frame(torch.zeros(2, 8), active) + codec.decode_frames(torch.zeros(2, 8, 3), active) + codec.decode_frames(torch.zeros(2, 8, 3), active=None) + + assert torch.equal(_active_args(codec._downsample.call_args_list), active[None]) + assert torch.equal(_active_args(codec._upsample.call_args_list), torch.stack((active, active, all_active))) + assert torch.equal(_active_args(codec.encoder_transformer.step.call_args_list), active[None]) + assert torch.equal( + _active_args(codec.decoder_transformer.step.call_args_list), + torch.stack((active, active, all_active)), + ) + + +@pytest.mark.cpu +def test_ring_mixed_offsets_match_singleton_streams() -> None: + batch_size, heads, head_dim, capacity, tokens = 3, 2, 3, 6, 2 + batched = _RingKV(batch_size, heads, head_dim, capacity, torch.device("cpu"), torch.float32) + singletons = [_RingKV(1, heads, head_dim, capacity, torch.device("cpu"), torch.float32) for _ in range(batch_size)] + reference_positions: list[torch.Tensor | None] = [None] * batch_size + active_schedule = [ + _mask(True, True, True), + _mask(True, False, False), + _mask(False, False, True), + _mask(True, True, False), + _mask(False, True, True), + _mask(True, False, True), + _mask(True, True, True), + ] * 2 + + torch.manual_seed(SEED) + for active in active_schedule: + keys = torch.randn(batch_size, heads, tokens, head_dim) + values = torch.randn_like(keys) + _, _, positions = batched.complete(keys, values, active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + _, _, singleton_positions = singletons[row].complete( + keys[row : row + 1], + values[row : row + 1], + _mask(True), + ) + reference_positions[row] = singleton_positions[0].clone() + + assert reference_positions[row] is not None + _assert_valid_ring_row_matches( + batched, + singletons[row], + positions, + reference_positions[row], + row, + ) + + assert any(int(offset) > capacity for offset in batched.end_offset) + assert not torch.equal(batched.end_offset[0], batched.end_offset[1]) + + +@pytest.mark.cpu +def test_ring_slot_recycle_masks_old_history() -> None: + ring = _RingKV(2, 1, 1, 8, torch.device("cpu"), torch.float32) + for _ in range(4): + ring.complete(torch.randn(2, 1, 2, 1), torch.randn(2, 1, 2, 1), _mask(True, True)) + + old_end = ring.end_offset.clone() + ring.reset_slot(1) + _, _, positions = ring.complete( + torch.randn(2, 1, 2, 1), + torch.randn(2, 1, 2, 1), + _mask(False, True), + ) + visible = positions[1] >= 0 + assert torch.all(positions[1][visible] >= old_end[1]) + + ring.reset_slot(0) + ring.bump_slot_start(0) + _, _, positions = ring.complete( + torch.randn(2, 1, 2, 1), + torch.randn(2, 1, 2, 1), + _mask(True, False), + ) + visible = positions[0] >= 0 + assert torch.all(positions[0][visible] >= old_end[0] + 1) + + +@pytest.mark.cpu +def test_temporal_streaming_rejects_invalid_active_shape() -> None: + temporal = _make_temporal(2, torch.device("cpu")) + + with pytest.raises(ValueError, match=r"active must have shape \(2,\)"): + temporal.step(torch.zeros(2, 1, 16), torch.ones(1)) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mimi_transformer_mixed_offsets_match_singletons() -> None: + batch_size, tokens, dim = 3, 2, 32 + batched = _make_mimi_transformer(batch_size, CUDA_DEVICE) + singletons = [_make_mimi_transformer(1, CUDA_DEVICE) for _ in range(batch_size)] + active_schedule = [ + _mask(True, True, True, device=CUDA_DEVICE), + _mask(True, False, False, device=CUDA_DEVICE), + _mask(False, False, True, device=CUDA_DEVICE), + _mask(True, True, False, device=CUDA_DEVICE), + _mask(False, True, True, device=CUDA_DEVICE), + _mask(True, False, True, device=CUDA_DEVICE), + _mask(True, True, True, device=CUDA_DEVICE), + ] * 2 + + for step, active in enumerate(active_schedule): + inputs = _inputs((batch_size, tokens, dim), SEED + step, CUDA_DEVICE) + offsets_before = batched._offset.clone() + cache_before = [kv.cache.clone() for kv in batched._kv] + output = batched.step(inputs, active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + expected = singletons[row].step( + inputs[row : row + 1], + _mask(True, device=CUDA_DEVICE), + ) + torch.testing.assert_close(output[row : row + 1], expected, rtol=0.0, atol=0.0) + assert torch.equal(batched._offset[row], singletons[row]._offset[0]) + for batched_kv, singleton_kv in zip(batched._kv, singletons[row]._kv): + assert torch.equal(batched_kv.end_offset[row], singleton_kv.end_offset[0]) + torch.testing.assert_close(batched_kv.cache[:, row], singleton_kv.cache[:, 0], rtol=0.0, atol=0.0) + else: + assert torch.equal(batched._offset[row], offsets_before[row]) + for kv, previous_cache in zip(batched._kv, cache_before): + assert torch.equal(kv.end_offset[row], offsets_before[row]) + assert torch.equal(kv.cache[:, row], previous_cache[:, row]) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_temporal_streaming_mixed_offsets_match_isolated_rows() -> None: + # Keep the reference shape at B=2 so exact comparisons use the same CUDA + # kernel shape while the target row runs independently. + batched = _make_temporal(2, CUDA_DEVICE) + references = [_make_temporal(2, CUDA_DEVICE) for _ in range(2)] + active_schedule = [ + _mask(True, True, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(False, False, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + ] * 2 + + for step, active in enumerate(active_schedule): + inputs = _inputs((2, 1, 16), SEED + 100 + step, CUDA_DEVICE) + offsets_before = batched._offset.clone() + cache_before = [kv.cache.clone() for kv in batched._kv] + output, logits = batched.step(inputs, active) + + for row, (reference, is_active) in enumerate(zip(references, active.tolist())): + if is_active: + reference_inputs = torch.zeros_like(inputs) + reference_inputs[row] = inputs[row] + reference_active = torch.zeros_like(active) + reference_active[row] = True + expected_output, expected_logits = reference.step(reference_inputs, reference_active) + torch.testing.assert_close(output[row : row + 1], expected_output[row : row + 1], rtol=0.0, atol=0.0) + torch.testing.assert_close(logits[row : row + 1], expected_logits[row : row + 1], rtol=0.0, atol=0.0) + assert torch.equal(batched._offset[row], reference._offset[row]) + for batched_kv, reference_kv in zip(batched._kv, reference._kv): + assert torch.equal(batched_kv.end_offset[row], reference_kv.end_offset[row]) + torch.testing.assert_close(batched_kv.cache[:, row], reference_kv.cache[:, row], rtol=0.0, atol=0.0) + else: + assert torch.equal(batched._offset[row], offsets_before[row]) + for kv, previous_cache in zip(batched._kv, cache_before): + assert torch.equal(kv.end_offset[row], offsets_before[row]) + assert torch.equal(kv.cache[:, row], previous_cache[:, row]) + + +@pytest.mark.parametrize("kind", ["conv", "convtr"]) +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mimi_conv_carries_preserve_inactive_rows(kind: str) -> None: + if kind == "conv": + conv = nn.Conv1d(1, 2, kernel_size=3, stride=2, device=CUDA_DEVICE) + batched = _StreamConv1d(conv, pad_mode="replicate") + row0 = _StreamConv1d(conv, pad_mode="replicate") + row1 = _StreamConv1d(conv, pad_mode="replicate") + samples = 4 + else: + conv = nn.ConvTranspose1d(1, 2, kernel_size=4, stride=2, device=CUDA_DEVICE) + batched = _StreamConvTr1d(conv) + row0 = _StreamConvTr1d(conv) + row1 = _StreamConvTr1d(conv) + samples = 2 + + for stream, batch_size in ((batched, 2), (row0, 1), (row1, 1)): + stream.reset(batch_size, CUDA_DEVICE, torch.float32) + + singletons = [row0, row1] + active_schedule = [ + _mask(False, False, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + ] + + for step, active in enumerate(active_schedule): + inputs = _inputs((2, 1, samples), SEED + 200 + step, CUDA_DEVICE) + carry_before = _stream_carry(batched).clone() + fresh_before = batched._fresh.clone() + output = PersonaPlexMimiCodec._run_stages(inputs, [(kind, batched)], active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + expected = PersonaPlexMimiCodec._run_stages( + inputs[row : row + 1], + [(kind, singletons[row])], + _mask(True, device=CUDA_DEVICE), + ) + torch.testing.assert_close(output[row : row + 1], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close( + _stream_carry(batched)[row], _stream_carry(singletons[row])[0], rtol=0.0, atol=0.0 + ) + assert torch.equal(batched._fresh[row], singletons[row]._fresh[0]) + else: + assert torch.equal(_stream_carry(batched)[row], carry_before[row]) + assert torch.equal(batched._fresh[row], fresh_before[row]) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 3b8813faf..356694380 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -43,6 +43,19 @@ FRAME_SIZE = 1920 CODEBOOKS = 8 +def _normalize_active(active: torch.Tensor | None, all_active: torch.Tensor) -> torch.Tensor: + # Stage 0 and Code2Wav currently use one B=1 codec per session and omit + # `active`, so preserve that API by treating None as all rows active. When + # they switch to shared B>1 codecs, their batch builders must pass bool[B]: + # True advances that row's streaming state; False keeps an absent or padded + # row's offsets and convolution carries unchanged. + if active is None: + return all_active + if active.shape != all_active.shape: + raise ValueError(f"active must have shape {tuple(all_active.shape)}, got {tuple(active.shape)}") + return active.to(device=all_active.device, dtype=torch.bool) + + def _map_moshi_codec_weights( state_dict: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]: @@ -107,23 +120,24 @@ class _StreamConv1d: def reset(self, batch_size: int, device, dtype) -> None: pad = self.kernel - self.stride self.prev = torch.zeros(batch_size, self.conv.in_channels, pad, device=device, dtype=dtype) - self._fresh = torch.ones(batch_size, dtype=torch.bool) + self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) def reset_slot(self, b: int) -> None: self.prev[b].zero_() self._fresh[b] = True - def __call__(self, x: torch.Tensor) -> torch.Tensor: - if self.pad_mode == "replicate" and bool(self._fresh.any()): + def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + if self.pad_mode == "replicate": pad = self.prev.shape[-1] edge = x[..., 0:1].expand(-1, -1, pad) - fresh = self._fresh.to(x.device).view(-1, 1, 1) + fresh = (self._fresh & active).view(-1, 1, 1) self.prev = torch.where(fresh, edge.to(self.prev.dtype), self.prev) - self._fresh[:] = False + self._fresh[active] = False x = torch.cat([self.prev, x], dim=-1) t = x.shape[-1] num_frames = max(0, (t - self.kernel) // self.stride + 1) - self.prev = x[..., num_frames * self.stride :] + prev = x[..., num_frames * self.stride :] + self.prev = torch.where(active.view(-1, 1, 1), prev, self.prev) if num_frames == 0: return x.new_zeros(x.shape[0], self.conv.out_channels, 0) return self.conv(x[..., : (num_frames - 1) * self.stride + self.kernel]) @@ -137,18 +151,19 @@ class _StreamConvTr1d: self.kernel = conv.kernel_size[0] self.stride = conv.stride[0] self.partial: torch.Tensor | None = None + self._fresh: torch.Tensor | None = None def reset(self, batch_size: int, device, dtype) -> None: self.partial = torch.zeros( batch_size, self.conv.out_channels, self.kernel - self.stride, device=device, dtype=dtype ) - self._fresh = torch.ones(batch_size, dtype=torch.bool) + self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) def reset_slot(self, b: int) -> None: self.partial[b].zero_() self._fresh[b] = True - def __call__(self, x: torch.Tensor) -> torch.Tensor: + def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: out = self.conv(x) length = out.shape[-1] tail = self.kernel - self.stride @@ -159,10 +174,11 @@ class _StreamConvTr1d: # it again, so subtract one copy -- except on a row's very first # frame, where the carry is zeros by construction. merge = merge - self.conv.bias[:, None] - merge[self._fresh] = 0.0 - self._fresh[:] = False + merge[self._fresh & active] = 0.0 + self._fresh[active] = False out[..., :pt] += merge - self.partial = out[..., length - tail :].clone() + partial = out[..., length - tail :].clone() + self.partial = torch.where(active.view(-1, 1, 1), partial, self.partial) return out[..., : length - tail] @@ -183,16 +199,23 @@ class _MimiTransformerLayer(nn.Module): self.scale1 = nn.Parameter(torch.empty(dim)) self.scale2 = nn.Parameter(torch.empty(dim)) - def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + kv: _RingKV, + offset: torch.Tensor, + context: int, + active: torch.Tensor, + ) -> torch.Tensor: B, T, _ = x.shape h = self.norm1(x) qkv = F.linear(h, self.in_proj_weight) qkv = qkv.view(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] q, k = _apply_rope(q, k, offset) - keys, values, pos_k = kv.complete(k, v) + keys, values, pos_k = kv.complete(k, v, active=active) pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1]) - pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) delta = pos_q - pos_k attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context) attn = F.scaled_dot_product_attention(q, keys, values, attn_bias.unsqueeze(1), dropout_p=0.0) @@ -218,7 +241,7 @@ class _MimiStreamingTransformer(nn.Module): heads = self.layers[0].num_heads hd = self.layers[0].head_dim self._kv = [_RingKV(batch_size, heads, hd, self.context, p.device, p.dtype) for _ in self.layers] - self._offset = torch.zeros(1, device=p.device, dtype=torch.long) + self._offset = torch.zeros(batch_size, device=p.device, dtype=torch.long) def reset_streaming(self) -> None: for kv in self._kv: @@ -229,11 +252,11 @@ class _MimiStreamingTransformer(nn.Module): for kv in self._kv: kv.reset_slot(b) - def step(self, x: torch.Tensor) -> torch.Tensor: + def step(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: """``x`` is ``[B, T, dim]`` (T = positions this frame, typically 2).""" for layer, kv in zip(self.layers, self._kv): - x = layer(x, kv, self._offset, self.context) - self._offset.add_(x.shape[1]) + x = layer(x, kv, self._offset, self.context, active) + self._offset.add_(x.shape[1] * active.to(self._offset.dtype)) return x def load_weights(self, state_dict: dict[str, torch.Tensor], prefix: str) -> int: @@ -339,6 +362,7 @@ class PersonaPlexMimiCodec(nn.Module): self._upsample = _StreamConvTr1d(m.upsample.conv) self._dec_stages = _walk_seanet(m.decoder.layers) self._batch_size: int | None = None + self._all_active: torch.Tensor # -- streaming state ------------------------------------------------------ @@ -354,6 +378,7 @@ class PersonaPlexMimiCodec(nn.Module): def streaming_init(self, batch_size: int) -> None: self._batch_size = batch_size + self._all_active = torch.ones(batch_size, dtype=torch.bool, device=self.device) for s in self._conv_states(): s.reset(batch_size, self.device, self.dtype) self.encoder_transformer.streaming_init(batch_size) @@ -376,39 +401,42 @@ class PersonaPlexMimiCodec(nn.Module): # -- per-frame codec ------------------------------------------------------- @staticmethod - def _run_stages(x: torch.Tensor, stages) -> torch.Tensor: + def _run_stages(x: torch.Tensor, stages, active: torch.Tensor) -> torch.Tensor: for kind, stage in stages: if kind == "res": act0, conv1, act2, conv3 = stage - x = x + conv3(act2(conv1(act0(x)))) + x = x + conv3(act2(conv1(act0(x), active)), active) else: - x = stage(x) + x = stage(x, active) if kind in {"conv", "convtr"} else stage(x) return x @torch.no_grad() - def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor: + def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, frame_size]`` float PCM -> ``[B, 8]`` codes.""" x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE) - x = self._run_stages(x, self._enc_stages) - x = self.encoder_transformer.step(x.transpose(1, 2)).transpose(1, 2) - x = self._downsample(x) + active = _normalize_active(active, self._all_active) + x = self._run_stages(x, self._enc_stages, active) + x = self.encoder_transformer.step(x.transpose(1, 2), active).transpose(1, 2) + x = self._downsample(x, active) codes = self.model.quantizer.encode(x) # [Q, B, T] return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous() @torch.no_grad() - def decode_frame(self, codes: torch.Tensor) -> torch.Tensor: + def decode_frame(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, 8]`` codes -> ``[B, frame_size]`` float PCM.""" emb = self.model.quantizer.decode(codes.to(self.device).view(-1, CODEBOOKS, 1)) - emb = self._upsample(emb) - emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2) - x = self._run_stages(emb, self._dec_stages) + active = _normalize_active(active, self._all_active) + emb = self._upsample(emb, active) + emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2) + x = self._run_stages(emb, self._dec_stages, active) return x[:, 0, :] @torch.no_grad() - def decode_frames(self, codes: torch.Tensor) -> torch.Tensor: + def decode_frames(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, 8, F]`` codes -> ``[B, F * frame_size]`` float PCM.""" emb = self.model.quantizer.decode(codes.to(self.device)) - emb = self._upsample(emb) - emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2) - x = self._run_stages(emb, self._dec_stages) + active = _normalize_active(active, self._all_active) + emb = self._upsample(emb, active) + emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2) + x = self._run_stages(emb, self._dec_stages, active) return x[:, 0, :] diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py index 717e4d505..28759dafb 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Streaming Helium temporal transformer for PersonaPlex (moshi-free, plain torch). Frame-clocked duplex needs a stateful per-frame forward: one 80 ms step consumes @@ -49,8 +49,8 @@ def _apply_rope(q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor, max_peri B, H, T, D = q.shape ds = torch.arange(D // 2, device=q.device, dtype=torch.float32) freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) - ts = offset.float() + torch.arange(T, device=q.device, dtype=torch.float32) - ts = ts.view(1, -1, 1) + ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + ts = ts.view(-1, 1, T, 1) dims = q.shape[:-1] q = q.view(*dims, D // 2, 2) @@ -80,7 +80,7 @@ class _RingKV: def __init__(self, batch_size: int, num_heads: int, dim_per_head: int, capacity: int, device, dtype): self.capacity = capacity self.cache = torch.zeros((2, batch_size, num_heads, capacity, dim_per_head), device=device, dtype=dtype) - self.end_offset = torch.zeros(1, device=device, dtype=torch.long) + self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long) self.start_offset = torch.zeros(batch_size, device=device, dtype=torch.long) def reset(self) -> None: @@ -90,23 +90,36 @@ class _RingKV: def reset_slot(self, b: int) -> None: # Mask everything written so far for row b; the row's next write is its # first visible entry. (LM sacrifice-tick +1 is applied by the caller.) - self.start_offset[b] = self.end_offset.clone() + self.start_offset[b] = self.end_offset[b] def bump_slot_start(self, b: int) -> None: self.start_offset[b] += 1 - def complete(self, k: torch.Tensor, v: torch.Tensor): + def complete(self, k: torch.Tensor, v: torch.Tensor, active: torch.Tensor): B, H, T, D = k.shape - indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + self.end_offset - indexes = indexes % self.capacity - self.cache[0].index_copy_(2, indexes, k) - self.cache[1].index_copy_(2, indexes, v) - self.end_offset.add_(T) + indexes = ( + torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype).view(1, -1) + + self.end_offset.view(-1, 1) + ) % self.capacity + idx4 = indexes.view(B, 1, T, 1).expand(-1, H, -1, D) + # Keep inactive rows completely inert. Once the ring is full, the + # physical future-write slot is also addressable by the position mask; + # writing it for an inactive row would therefore leak padded data into + # its next attention step. + active_view = active.view(B, 1, 1, 1) + old_k = self.cache[0].gather(2, idx4) + old_v = self.cache[1].gather(2, idx4) + k = torch.where(active_view, k, old_k) + v = torch.where(active_view, v, old_v) + self.cache[0].scatter_(2, idx4, k) + self.cache[1].scatter_(2, idx4, v) + self.end_offset.add_(T * active.to(self.end_offset.dtype)) idx = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long) - invalid = idx >= self.end_offset - end_index = self.end_offset % self.capacity - delta = idx - end_index + end_offset = self.end_offset.view(-1, 1) + invalid = idx.view(1, -1) >= end_offset + end_index = end_offset % self.capacity + delta = idx.view(1, -1) - end_index # `delta <= 0` (not `< 0`) is moshi's exact convention (transformer.py # RingKVCache.complete). It labels the just-past-newest slot as the future # write position, so once the ring has wrapped the single oldest in-window @@ -114,14 +127,19 @@ class _RingKV: # verbatim from the reference and only shows after the window fills (Helium # ~3000 frames / 240 s); it costs one frame out of thousands. Do NOT change # this to `< 0`: it would diverge from moshi and break greedy bit-parity. - positions = torch.where(delta <= 0, self.end_offset + delta, self.end_offset + delta - self.capacity) + positions = torch.where(delta <= 0, end_offset + delta, end_offset + delta - self.capacity) positions = torch.where(invalid, torch.full_like(positions, -1), positions) - positions = positions.view(1, -1) # [1, capacity] below = positions < self.start_offset.view(-1, 1) # [B, capacity] positions = torch.where(below, torch.full_like(positions, -1), positions) return self.cache[0], self.cache[1], positions +def _normalize_temporal_active(active: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + if active.shape != reference.shape: + raise ValueError(f"active must have shape {tuple(reference.shape)}, got {tuple(active.shape)}") + return active.to(device=reference.device, dtype=torch.bool) + + class _TemporalLayer(nn.Module): """One Helium decoder layer in the checkpoint's fused layout.""" @@ -137,7 +155,14 @@ class _TemporalLayer(nn.Module): self.norm1_alpha = nn.Parameter(torch.ones(1, 1, dim)) self.norm2_alpha = nn.Parameter(torch.ones(1, 1, dim)) - def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + kv: _RingKV, + offset: torch.Tensor, + context: int, + active: torch.Tensor, + ) -> torch.Tensor: B, T, _ = x.shape h = _rms_norm_f32(x, self.norm1_alpha, 1e-8) qkv = F.linear(h, self.in_proj_weight) @@ -147,9 +172,9 @@ class _TemporalLayer(nn.Module): q, k, v = qkv[0], qkv[1], qkv[2] q, k = _apply_rope(q, k, offset) - keys, values, pos_k = kv.complete(k, v) + keys, values, pos_k = kv.complete(k, v, active) pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1]) # [B, 1, cap] - pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) delta = pos_q - pos_k attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context) attn_bias = attn_bias.unsqueeze(1) # [B, 1, T, cap] @@ -199,7 +224,7 @@ class PersonaPlexTemporalStreaming(nn.Module): head_dim = self.layers[0].head_dim # Capacity = context window (the mask truncates at `context` anyway). self._kv = [_RingKV(batch_size, heads, head_dim, self.context, p.device, p.dtype) for _ in self.layers] - self._offset = torch.zeros(1, device=p.device, dtype=torch.long) + self._offset = torch.zeros(batch_size, device=p.device, dtype=torch.long) def reset_streaming(self) -> None: assert self._kv is not None, "call streaming_init first" @@ -221,12 +246,17 @@ class PersonaPlexTemporalStreaming(nn.Module): # -- per-frame step ------------------------------------------------------- @torch.no_grad() - def step(self, frame_embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def step( + self, + frame_embedding: torch.Tensor, + active: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: assert self._kv is not None, "call streaming_init first" + active = _normalize_temporal_active(active, self._offset) x = frame_embedding for layer, kv in zip(self.layers, self._kv): - x = layer(x, kv, self._offset, self.context) - self._offset.add_(x.shape[1]) + x = layer(x, kv, self._offset, self.context, active) + self._offset.add_(x.shape[1] * active.to(self._offset.dtype)) out = _rms_norm_f32(x, self.out_norm_alpha, 1e-8) text_logits = F.linear(out, self.text_linear) return out, text_logits[:, None] From e3be42e052e88d052d21cf7a24ee84c70f020fd6 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:33:40 +0800 Subject: [PATCH 27/99] [Core] Give the OpenAI Realtime wire codec a reusable home (RFC #6592 P0a) (#7640) Signed-off-by: psv666 <2693925048@qq.com> --- .../common/ci_source_file_dependencies.yml | 2 + .github/CODEOWNERS | 5 + docs/design/fullduplex.md | 77 +- docs/design/module/entrypoints.md | 43 + tests/protocol/__init__.py | 2 + tests/protocol/duplex/__init__.py | 2 + .../duplex/test_duplex_protocol_facade.py | 144 ++++ tests/protocol/realtime/__init__.py | 2 + .../realtime/test_protocol_import_boundary.py | 106 +++ .../realtime/test_realtime_capabilities.py | 91 +++ .../test_realtime_codec_single_source.py | 256 ++++++ tools/pre_commit/check_forbidden_imports.py | 7 +- vllm_omni/clients/inline_duplex.py | 4 +- vllm_omni/engine/duplex/audio.py | 250 +----- vllm_omni/engine/duplex/commands.py | 238 +++--- vllm_omni/engine/duplex/config.py | 36 +- vllm_omni/engine/duplex/events.py | 737 ++---------------- vllm_omni/engine/duplex/realtime_commands.py | 663 ++-------------- vllm_omni/engine/duplex/realtime_events.py | 12 +- vllm_omni/engine/duplex/session/runner.py | 2 +- .../entrypoints/duplex/realtime_input.py | 10 +- vllm_omni/entrypoints/duplex/serving.py | 2 +- vllm_omni/protocol/__init__.py | 17 + vllm_omni/protocol/duplex/__init__.py | 159 ++++ vllm_omni/protocol/duplex/commands.py | 164 ++++ vllm_omni/protocol/duplex/errors.py | 75 ++ vllm_omni/protocol/duplex/events.py | 419 ++++++++++ vllm_omni/protocol/realtime/__init__.py | 163 ++++ vllm_omni/protocol/realtime/audio.py | 218 ++++++ vllm_omni/protocol/realtime/audio_input.py | 243 ++++++ vllm_omni/protocol/realtime/capabilities.py | 105 +++ vllm_omni/protocol/realtime/commands.py | 126 +++ vllm_omni/protocol/realtime/errors.py | 37 + vllm_omni/protocol/realtime/events.py | 467 +++++++++++ vllm_omni/protocol/realtime/formats.py | 177 +++++ vllm_omni/protocol/realtime/items.py | 197 +++++ vllm_omni/protocol/realtime/session.py | 163 ++++ 37 files changed, 3792 insertions(+), 1629 deletions(-) create mode 100644 tests/protocol/__init__.py create mode 100644 tests/protocol/duplex/__init__.py create mode 100644 tests/protocol/duplex/test_duplex_protocol_facade.py create mode 100644 tests/protocol/realtime/__init__.py create mode 100644 tests/protocol/realtime/test_protocol_import_boundary.py create mode 100644 tests/protocol/realtime/test_realtime_capabilities.py create mode 100644 tests/protocol/realtime/test_realtime_codec_single_source.py create mode 100644 vllm_omni/protocol/__init__.py create mode 100644 vllm_omni/protocol/duplex/__init__.py create mode 100644 vllm_omni/protocol/duplex/commands.py create mode 100644 vllm_omni/protocol/duplex/errors.py create mode 100644 vllm_omni/protocol/duplex/events.py create mode 100644 vllm_omni/protocol/realtime/__init__.py create mode 100644 vllm_omni/protocol/realtime/audio.py create mode 100644 vllm_omni/protocol/realtime/audio_input.py create mode 100644 vllm_omni/protocol/realtime/capabilities.py create mode 100644 vllm_omni/protocol/realtime/commands.py create mode 100644 vllm_omni/protocol/realtime/errors.py create mode 100644 vllm_omni/protocol/realtime/events.py create mode 100644 vllm_omni/protocol/realtime/formats.py create mode 100644 vllm_omni/protocol/realtime/items.py create mode 100644 vllm_omni/protocol/realtime/session.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index e29001d9a..c101db48d 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -44,6 +44,7 @@ x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/entrypoints/duplex/ - vllm_omni/entrypoints/duplex_omni.py - vllm_omni/entrypoints/async_omni_base.py + - vllm_omni/protocol/ - vllm_omni/clients/ - vllm_omni/outputs/duplex.py - vllm_omni/model_executor/duplex_sampling.py @@ -259,6 +260,7 @@ source_file_dependencies: - tests/entrypoints/test_duplex_omni.py - tests/entrypoints/openai_api/test_duplex_api_server.py - tests/model_executor/models/minicpmo_4_5/duplex/ + - tests/protocol/ omni_minicpmo_4_5_duplex_perf: - *minicpmo_4_5 diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index c92c7b22f..e77f70914 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -67,6 +67,11 @@ # File-level overrides within entrypoints/ /vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao /vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter +# Shared wire codec: protocol/realtime (OpenAI Realtime GA) and protocol/duplex +# (vLLM-Omni's extension of it) — RFC #6592 P0a. Layer-neutral by design — the +# engine depends on it — so it is owned by the serving-boundary document rather +# than by any one consumer. — docs/design/module/entrypoints.md +/vllm_omni/protocol/ @alex-jw-brooks @linyueqian @NickCao # Error contracts — docs/design/module/error_contracts.md /vllm_omni/errors.py @alex-jw-brooks @NickCao diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 41a050a47..324cd5bdc 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -9,6 +9,12 @@ the owner document for `vllm_omni/engine/duplex/**`, `vllm_omni/entrypoints/duplex_omni.py`, `vllm_omni/entrypoints/duplex/**`, `vllm_omni/clients/**` and the per-model `duplex/` packages. +`vllm_omni/protocol/realtime/**` (the OpenAI Realtime API) and +`vllm_omni/protocol/duplex/**` (vLLM-Omni's extension of it) are the shared +wire codec duplex is built on. Duplex is its first consumer but not its owner: the codec +is deliberately runtime-agnostic so a second Realtime surface can reuse it +(RFC #6592 P0a). + The public wire contract lives in [`docs/serving/realtime_duplex_api.md`](../serving/realtime_duplex_api.md); this page describes how it is produced. @@ -124,6 +130,40 @@ plugin and the session runtime config to `DuplexOrchestrator` directly. attachment/resume/replay bookkeeping; it holds no session state. 6. **No `typing.Protocol`** in the duplex surfaces: the plugin, data plane, session state, PCM buffer, stage port and client transport seams are ABCs. +7. **The wire codec is not duplex, and a wire type is not a mailbox message.** + The event and command *vocabulary* plus its wire rendering, the audio format + negotiation, the conversation-item rules and the error envelope are model- + and runtime-agnostic, and split along the three tiers + [`docs/serving/realtime_duplex_api.md`](../serving/realtime_duplex_api.md) + already defines: **Tier 1** (identical to OpenAI) is + `vllm_omni/protocol/realtime/**`; **Tier 2** (OpenAI names carrying our + extensions) and **Tier 3** (ours alone) are `vllm_omni/protocol/duplex/**`. + A Tier 2 class *subclasses* its Tier 1 twin and declares only what it adds, + so the OpenAI surface stays honestly pure --- `session.created` without + `resume_token`, `input_audio_buffer.append` without `video_frames`. Our + error **codes** are Tier 3 (`protocol/duplex/errors.py`); only the envelope + shape and OpenAI's three `error.type` classes are Tier 1. + What stays engine-side is everything that is *internal representation* + rather than contract: `DuplexCommand.payload()` renders the session runner's + mailbox dictionary, whose channel genuinely differs from the client event + (`session.update` and the three `conversation.item.*` commands all travel on + `turn.signal`); `engine/duplex/realtime_commands.py` decides which command a + decoded event becomes; `engine/duplex/realtime_events.py` holds the + session's projection state. The codec may not import the engine, the + entrypoints, `model_executor` or the clients + (`tests/protocol/realtime/test_protocol_import_boundary.py`), and there is + exactly one implementation of each codec behaviour + (`tests/protocol/realtime/test_realtime_codec_single_source.py`). + The dependency chain is `protocol/realtime` <- `protocol/duplex` <- engine / + entrypoints / clients, and a duplex consumer uses **only** the middle link: + `protocol/duplex/**` re-exports the Tier 1 names it does not extend, so a + helper that later needs a duplex-specific version (`convert_input_audio_with_rate` + resamples to MiniCPM-o's 16 kHz rather than the client's rate) is overridden + in one file instead of at every call site + (`tests/protocol/duplex/test_duplex_protocol_facade.py`). + `DuplexEvent` is therefore a plain alias of `RealtimeEvent` --- an event has + no engine-internal half --- while `DuplexCommand` is a real class, because a + command does. ## Package layout @@ -142,6 +182,31 @@ vllm_omni/ │ │ ├── chat_completions.py DuplexChatCompletionsAdapter (/v1/chat/completions on a session per request) │ │ └── websocket.py websocket send/close/receive helpers │ └── openai/api_server.py builds DuplexOmni for duplex models; session-backed app state +├── protocol/ SHARED WIRE CODEC (no engine / no model / no transport) +│ ├── realtime/ the OpenAI Realtime API itself +│ │ ├── events.py RealtimeEvent + OpenAI's 30 server events (22 ship as Tier 1; +│ │ │ 8 have a Tier 2 twin in protocol/duplex) +│ │ ├── commands.py RealtimeCommand + OpenAI's 10 client commands (8 Tier 1, +│ │ │ 2 with a Tier 2 twin) --- wire_type + fields, no payload() +│ │ ├── formats.py audio format spellings, what is supported, per-session validation +│ │ ├── session.py session-object readers; RealtimeInputDefaults (append defaults) +│ │ ├── items.py conversation item shape, truncation, transcripts, camera frames +│ │ ├── audio_input.py decode_audio_append -> RealtimeAudioAppend; client speech hints +│ │ ├── audio.py PCM / G.711 / WAV decode, resample, re-encode +│ │ ├── errors.py RealtimeProtocolError (the envelope's exception type) +│ │ └── capabilities.py RealtimeProtocolCapabilities: what one consumer can serve +│ └── duplex/ Tier 2 + Tier 3, and the one door a duplex consumer uses +│ ├── __init__.py re-exports the Tier 1 helper functions (audio, formats, +│ │ items, session, capabilities) so the engine never +│ │ imports protocol/realtime directly +│ ├── events.py the whole 42-event vocabulary: 22 Tier-1 re-exports, +│ │ 8 Tier-2 subclasses (session.created + resume_token, ...), +│ │ 12 Tier-3 events (listen/speak, playback, resume, ...); +│ │ DuplexEvent = RealtimeEvent +│ ├── commands.py the whole 17-command vocabulary: 8 Tier-1 re-exports, +│ │ 2 Tier-2 subclasses (append + hints, commit + final), +│ │ 7 Tier-3 commands (barge_in, playback.ack, ...) +│ └── errors.py our error-code vocabulary (Tier 3) + RealtimeProtocolError ├── engine/ │ ├── omni_engine_base.py OmniEngineBase (stage processes, orchestrator thread, queues, RPC) │ ├── async_omni_engine.py AsyncOmniEngine (turn-based request building) @@ -149,16 +214,20 @@ vllm_omni/ │ ├── orchestrator.py OrchestratorBase + Orchestrator │ ├── duplex_orchestrator.py DuplexOrchestrator (+ DuplexOrchestratorRequestState; implements DuplexStagePort) │ └── duplex/ -│ ├── commands.py DuplexCommand dataclasses, command_from_realtime -│ ├── realtime_commands.py Realtime client event -> DuplexCommand translation -│ ├── events.py DuplexEvent dataclasses (+ to_realtime), REALTIME_ERROR_TYPES_BY_CODE +│ ├── commands.py mailbox half: DuplexCommand.payload() + the mailbox `type`, +│ │ paired with each protocol command; command_from_realtime +│ ├── realtime_commands.py duplex binding of the codec: decoded event -> DuplexCommand, +│ │ DUPLEX_REALTIME_CAPABILITIES, duplex_response_format +│ ├── events.py re-export shim (DuplexEvent = RealtimeEvent) + the runner's +│ │ epoch-filter sets DOMAIN_TERMINAL_EVENTS / MODEL_OUTPUT_EVENTS │ ├── realtime_events.py RealtimeProjectionState: internal event -> typed events │ ├── messages.py queue envelopes (Open/Close/Resume/Touch/Command/Result/Event), DuplexSessionError │ ├── config.py DuplexSessionConfig, DuplexCapabilities, ResponseCreateOptions │ ├── contracts.py DuplexFence (session_id, epoch, turn_id), stage request records, DuplexStagePort │ ├── plugin.py DuplexModelPlugin, DuplexModelSessionState, DuplexDataPlane, PcmAppendBuffer ABCs │ ├── turn_detection.py server-side VAD turn detector used by the session -│ ├── audio.py / vad.py / intermediate.py +│ ├── audio.py compatibility re-export, via protocol/duplex +│ ├── vad.py / intermediate.py │ └── session/ one engine-resident session and everything that runs it │ ├── engine_session.py DuplexEngineSession: ledgers, lease, fence, stage resources, append sequencing │ ├── runner.py DuplexSessionRunner (per-session mailbox on the orchestrator loop) diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index f37ddc6ba..0c0ba8439 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -30,6 +30,7 @@ primary_code_paths: - vllm_omni/entrypoints/utils.py - vllm_omni/entrypoints/duplex/** - vllm_omni/clients/** + - vllm_omni/protocol/** primary_path_exceptions: - path: vllm_omni/entrypoints/openai/errors.py owner: error_contracts.md @@ -53,6 +54,7 @@ depends_on: - error_contracts.md - engine_orchestration.md validation_paths: + - tests/protocol/** - tests/entrypoints/test_omni_entrypoints.py - tests/entrypoints/test_async_omni.py - tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -106,6 +108,31 @@ placement, payload implementation, or semantic error classification. `entrypoints/openai/errors.py` is an explicit primary-path exception owned by `error_contracts.md`. +It also owns `vllm_omni/protocol/**`, the shared wire codec promoted by RFC +[#6592](https://github.com/vllm-project/vllm-omni/issues/6592) P0a. It follows the three tiers +`docs/serving/realtime_duplex_api.md` already defines. `protocol/realtime/**` +is Tier 1, identical to OpenAI: 30 server events, 10 client commands, their +typed fields and pure wire rendering, audio-format negotiation, +conversation-item rules and the error-envelope shape. `protocol/duplex/**` is +Tier 2 (OpenAI names carrying vLLM-Omni extensions, each a subclass of its +Tier 1 twin declaring only what it adds) and Tier 3 (ours alone, including our +error-code vocabulary), kept separate so a GA-only consumer is not handed a +vocabulary its clients never send. That package sits +outside `entrypoints/` on purpose --- the engine depends on it, so it belongs +to neither layer --- but deciding what a Realtime client may put on the wire is +this document's subject, and the same reasoning already places +`vllm_omni/errors.py`, `inputs/` and `outputs/` under contract documents rather +than under a consumer. + +It is deliberately **not** owned by `../fullduplex.md`, even though duplex is +the codec's only consumer today. RFC #6592 asks for a named owner precisely so +a second Realtime consumer does not need duplex review to change the codec +(#6592 open question 8). Ownership here is review scope, not a dependency: the +codec may not import the engine, the entrypoints or the model code, which +`tests/protocol/realtime/test_protocol_import_boundary.py` enforces. When a +second consumer lands (#6592 P0b) the package is expected to graduate to its +own module document with owners drawn from both consumers. + ## Candidate invariants These identifiers are proposals while the document is `draft`. @@ -115,6 +142,22 @@ These identifiers are proposals while the document is `draft`. **Rule:** Entrypoints MUST NOT implement cross-stage routing or stage lifecycle policy. +### ENTRY-INV-002: The shared Realtime codec does not depend on its consumers + +**Rule:** `vllm_omni/protocol/**` MUST NOT import `vllm_omni.engine`, +`vllm_omni.entrypoints`, `vllm_omni.model_executor`, `vllm_omni.worker` or +`vllm_omni.clients`; `protocol/duplex/**` MAY depend on `protocol/realtime/**` +but never the reverse, and a Tier 1 class MUST NOT carry a vLLM-Omni extension +field; and each protocol/codec behavior MUST have exactly one +implementation that every consumer uses. A wire type MUST NOT carry a +consumer's internal representation: the duplex mailbox rendering +(`DuplexCommand.payload()`, whose channel differs from the client event for +`session.update` and the `conversation.item.*` commands) stays engine-side. +A duplex consumer (`engine/duplex/**`, `entrypoints/duplex/**`, the duplex +clients) MUST import `vllm_omni.protocol.duplex` and MUST NOT import +`vllm_omni.protocol.realtime` directly, so the tier boundary has exactly one +extension point. + ### ENTRY-INV-100: Public requests are normalized once **Rule:** Public protocol values MUST be validated and converted to internal diff --git a/tests/protocol/__init__.py b/tests/protocol/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/protocol/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/protocol/duplex/__init__.py b/tests/protocol/duplex/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/protocol/duplex/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/protocol/duplex/test_duplex_protocol_facade.py b/tests/protocol/duplex/test_duplex_protocol_facade.py new file mode 100644 index 000000000..61dd15fe1 --- /dev/null +++ b/tests/protocol/duplex/test_duplex_protocol_facade.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``protocol.duplex`` is the only door a duplex consumer uses. + +The layering is a chain:: + + protocol.realtime <- protocol.duplex <- engine / entrypoints / clients + +The middle link only earns its keep if the last arrow is the *only* one. If the +engine may also reach past it straight into ``protocol.realtime``, then giving a +Tier 1 helper a duplex-specific version later means touching every call site +instead of one file --- which is exactly the change +``convert_input_audio_with_rate`` is expected to need (it resamples to +MiniCPM-o's 16 kHz, not the client's rate). + +So this asserts two things: nobody skips the middle link, and the middle link is +complete enough that nobody has to. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[3] + +#: Every module that speaks the duplex protocol. These import +#: ``vllm_omni.protocol.duplex``; reaching into ``vllm_omni.protocol.realtime`` +#: from here skips the extension point. +CONSUMER_PATHS = ( + "vllm_omni/engine/duplex", + "vllm_omni/engine/duplex_omni_engine.py", + "vllm_omni/entrypoints/duplex", + "vllm_omni/entrypoints/duplex_omni.py", + "vllm_omni/clients/duplex.py", + "vllm_omni/clients/inline_duplex.py", +) + + +def _consumer_modules() -> list[Path]: + found: list[Path] = [] + for entry in CONSUMER_PATHS: + path = REPO_ROOT / entry + found.extend(sorted(path.rglob("*.py")) if path.is_dir() else [path]) + assert found, "consumer list is stale" + return found + + +@pytest.mark.parametrize("module", _consumer_modules(), ids=lambda p: p.name) +def test_a_duplex_consumer_does_not_skip_the_duplex_protocol_package(module: Path) -> None: + # Source-level, so a lazy import inside a function body is caught too. + tree = ast.parse(module.read_text()) + offenders = [] + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and (node.module or "").startswith("vllm_omni.protocol.realtime"): + offenders.append(f"line {node.lineno}: from {node.module} import ...") + elif isinstance(node, ast.Import): + offenders.extend( + f"line {node.lineno}: import {alias.name}" + for alias in node.names + if alias.name.startswith("vllm_omni.protocol.realtime") + ) + assert not offenders, ( + f"{module.relative_to(REPO_ROOT)} imports the Tier 1 package directly; " + f"import from vllm_omni.protocol.duplex instead:\n " + "\n ".join(offenders) + ) + + +def test_the_duplex_package_carries_the_whole_command_vocabulary() -> None: + from vllm_omni.engine.duplex import commands as engine_commands + from vllm_omni.protocol.duplex import commands as duplex_wire + + # Every command the engine handles is nameable through the one door. + engine_names = { + name + for name in engine_commands.__all__ + if isinstance(getattr(engine_commands, name), type) + and issubclass(getattr(engine_commands, name), engine_commands.DuplexCommand) + and getattr(engine_commands, name) is not engine_commands.DuplexCommand + } + assert engine_names <= set(duplex_wire.__all__) + + +def test_the_duplex_package_carries_the_whole_event_vocabulary() -> None: + from vllm_omni.engine.duplex import events as engine_events + from vllm_omni.protocol.duplex import events as duplex_wire + + engine_names = { + name + for name in engine_events.__all__ + if isinstance(getattr(engine_events, name), type) + and issubclass(getattr(engine_events, name), engine_events.DuplexEvent) + } + assert engine_names <= set(duplex_wire.__all__) + + +@pytest.mark.parametrize( + "module_pair", + [("events", "events"), ("commands", "commands")], + ids=["events", "commands"], +) +def test_a_reexported_tier1_name_is_the_tier1_object(module_pair: tuple[str, str]) -> None: + """Re-export, not re-declaration: the Tier 1 class is one object.""" + import importlib + + duplex_mod = importlib.import_module(f"vllm_omni.protocol.duplex.{module_pair[0]}") + realtime_mod = importlib.import_module(f"vllm_omni.protocol.realtime.{module_pair[1]}") + + shared = [n for n in duplex_mod.__all__ if hasattr(realtime_mod, n)] + assert shared, "nothing re-exported --- the facade is not doing its job" + for name in shared: + duplex_obj = getattr(duplex_mod, name) + realtime_obj = getattr(realtime_mod, name) + if duplex_obj is realtime_obj: + continue + # The only permitted divergence is a Tier 2 class extending its twin. + assert isinstance(duplex_obj, type) and issubclass(duplex_obj, realtime_obj), ( + f"{name} in protocol.duplex is neither the Tier 1 object nor a subclass of it" + ) + + +def test_duplex_event_is_an_alias_but_duplex_command_is_not() -> None: + """The asymmetry is deliberate, and a later reader should not 'fix' it. + + An event has no engine-internal half --- what the session emits is what the + socket sends --- so ``DuplexEvent`` is simply ``RealtimeEvent``. A command + does: ``payload()`` renders the session runner's mailbox dictionary, and for + four of the seventeen the mailbox channel is not even the client event type. + """ + from vllm_omni.engine.duplex.commands import DuplexCommand + from vllm_omni.protocol.duplex.events import DuplexEvent + from vllm_omni.protocol.realtime.commands import RealtimeCommand + from vllm_omni.protocol.realtime.events import RealtimeEvent + + assert DuplexEvent is RealtimeEvent + + assert DuplexCommand is not RealtimeCommand + assert issubclass(DuplexCommand, RealtimeCommand) + assert hasattr(DuplexCommand, "payload") and not hasattr(RealtimeCommand, "payload") diff --git a/tests/protocol/realtime/__init__.py b/tests/protocol/realtime/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/protocol/realtime/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/protocol/realtime/test_protocol_import_boundary.py b/tests/protocol/realtime/test_protocol_import_boundary.py new file mode 100644 index 000000000..a00d4da62 --- /dev/null +++ b/tests/protocol/realtime/test_protocol_import_boundary.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The shared Realtime codec must not depend on any runtime that uses it. + +``vllm_omni.protocol.realtime`` exists so a second Realtime surface can reuse +the wire codec without adopting the duplex session control plane. That only +holds while the arrows point one way: the engine, the entrypoints and the model +code import the protocol, never the reverse. Nothing in the package's own code +enforces that, so it is asserted here --- once on the source (which catches a +lazy import inside a function body) and once on a real import (which catches a +dependency arriving through a third module). +""" + +from __future__ import annotations + +import ast +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[3] +PROTOCOL_DIR = REPO_ROOT / "vllm_omni" / "protocol" + +#: Packages that consume the protocol. A dependency on any of them would make +#: the codec unusable outside the runtime it accidentally reached into. +#: ``vllm_omni.utils`` is deliberately absent: it is leaf helpers (the audio +#: resampler) that consume nothing, so depending on it costs a reuser nothing. +FORBIDDEN_PREFIXES = ( + "vllm_omni.engine", + "vllm_omni.entrypoints", + "vllm_omni.model_executor", + "vllm_omni.worker", + "vllm_omni.clients", +) + + +def _protocol_sources() -> list[Path]: + sources = sorted(PROTOCOL_DIR.rglob("*.py")) + assert sources, f"no python sources under {PROTOCOL_DIR}" + return sources + + +def _imported_modules(source: Path) -> set[str]: + """Every module name imported anywhere in the file, lazy imports included.""" + names: set[str] = set() + for node in ast.walk(ast.parse(source.read_text())): + if isinstance(node, ast.Import): + names.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + names.add(node.module) + return names + + +@pytest.mark.parametrize("source", _protocol_sources(), ids=lambda path: path.name) +def test_protocol_sources_do_not_import_a_consumer(source: Path) -> None: + offenders = sorted( + name + for name in _imported_modules(source) + if any(name == prefix or name.startswith(prefix + ".") for prefix in FORBIDDEN_PREFIXES) + ) + assert not offenders, f"{source.relative_to(REPO_ROOT)} imports {', '.join(offenders)}" + + +def test_importing_the_protocol_package_loads_no_consumer() -> None: + # A child interpreter, because this test session has already imported the + # duplex engine: only a fresh process can tell what the package pulls in. + # ``vllm_omni`` itself is imported first, because its ``__init__`` loads a + # baseline graph (patches, the model registry) that has nothing to do with + # this package; only what the protocol import *adds* is measured. + env = os.environ.copy() + env.pop("CUDA_VISIBLE_DEVICES", None) + env.pop("HIP_VISIBLE_DEVICES", None) + script = f""" +import sys + +import vllm_omni # noqa: F401 + +baseline = set(sys.modules) + +import vllm_omni.protocol.realtime # noqa: F401 + +forbidden_prefixes = {FORBIDDEN_PREFIXES!r} +added = sorted( + name + for name in set(sys.modules) - baseline + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes) +) +if added: + raise SystemExit("the realtime protocol package loaded consumers: " + ", ".join(added)) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=180, + check=False, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/protocol/realtime/test_realtime_capabilities.py b/tests/protocol/realtime/test_realtime_capabilities.py new file mode 100644 index 000000000..2a05fc0a9 --- /dev/null +++ b/tests/protocol/realtime/test_realtime_capabilities.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``RealtimeProtocolCapabilities``: one consumer's answers, the codec's checking. + +This is the seam that lets a non-duplex Realtime surface validate a +``session.update`` without a ``DuplexModelPlugin``. What it has to get right is +narrow: the order the checks run in (formats before turn detection, because +that is the order clients already see), the error codes, and that a consumer +which declares a narrower format set is actually narrowed. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import pytest + +from vllm_omni.protocol.realtime.capabilities import ( + RealtimeProtocolCapabilities, + RealtimeSessionRejection, + validate_session_payload, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEFAULT = RealtimeProtocolCapabilities() + + +def test_a_consumer_without_a_turn_detection_check_accepts_a_plain_session(): + assert ( + validate_session_payload({"model": "m", "audio": {"input": {"format": "pcm16"}}}, capabilities=_DEFAULT) is None + ) + + +def test_a_consumer_without_a_turn_detection_check_ignores_turn_detection(): + payload = {"turn_detection": {"type": "semantic_vad"}} + + assert validate_session_payload(payload, capabilities=_DEFAULT) is None + + +def test_an_unsupported_input_format_is_rejected_with_the_realtime_code(): + rejection = validate_session_payload({"input_audio_format": "opus"}, capabilities=_DEFAULT) + + assert rejection == RealtimeSessionRejection( + code="unsupported_audio_format", message="Unsupported input_audio_format: opus" + ) + + +def test_a_consumer_may_declare_a_narrower_format_set(): + pcm_only = RealtimeProtocolCapabilities(input_audio_formats=frozenset({"pcm16"})) + + assert validate_session_payload({"input_audio_format": "pcm16"}, capabilities=pcm_only) is None + rejection = validate_session_payload({"input_audio_format": "g711_ulaw"}, capabilities=pcm_only) + assert rejection is not None + assert rejection.code == "unsupported_audio_format" + # The default capabilities still accept it: the narrowing is per consumer. + assert validate_session_payload({"input_audio_format": "g711_ulaw"}, capabilities=_DEFAULT) is None + + +def test_the_turn_detection_check_runs_and_reports_its_own_code(): + def reject_everything(session_payload: Mapping[str, object]) -> str | None: + del session_payload + return "server_vad is not available" + + capabilities = RealtimeProtocolCapabilities(validate_turn_detection=reject_everything) + + rejection = validate_session_payload({"turn_detection": {"type": "server_vad"}}, capabilities=capabilities) + + assert rejection == RealtimeSessionRejection( + code="unsupported_turn_detection", message="server_vad is not available", param="turn_detection" + ) + + +def test_a_bad_format_is_reported_before_a_bad_turn_detection(): + calls: list[Mapping[str, object]] = [] + + def record(session_payload: Mapping[str, object]) -> str | None: + calls.append(session_payload) + return "turn detection is wrong too" + + capabilities = RealtimeProtocolCapabilities(validate_turn_detection=record) + + rejection = validate_session_payload( + {"input_audio_format": "opus", "turn_detection": {"type": "semantic_vad"}}, + capabilities=capabilities, + ) + + assert rejection is not None + assert rejection.code == "unsupported_audio_format" + assert calls == [] diff --git a/tests/protocol/realtime/test_realtime_codec_single_source.py b/tests/protocol/realtime/test_realtime_codec_single_source.py new file mode 100644 index 000000000..2f719750b --- /dev/null +++ b/tests/protocol/realtime/test_realtime_codec_single_source.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""There is one Realtime codec, and the duplex engine uses it. + +RFC #6592 P0a rejects the "copy the codec and fix it up" option because the two +copies drift. The extraction is only worth anything while the duplex names are +*the same objects* as the protocol ones, not lookalikes --- so assert identity, +which a re-implementation cannot satisfy. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from vllm_omni.engine.duplex import audio as duplex_audio +from vllm_omni.engine.duplex import events as duplex_events +from vllm_omni.engine.duplex import realtime_commands as duplex_codec +from vllm_omni.engine.duplex.commands import DuplexCommandError +from vllm_omni.protocol import realtime as protocol +from vllm_omni.protocol.realtime.errors import RealtimeProtocolError + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Names ``engine.duplex.realtime_commands`` re-exports for its existing importers. +_REEXPORTED_FROM_PROTOCOL = ( + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_INPUT_HINT_KEYS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "RealtimeInputDefaults", + "apply_realtime_session_defaults", + "copy_realtime_input_hints", + "input_audio_transcription_config", + "input_explicitly_non_speech", + "input_looks_like_speech", + "input_transcript_from_item", + "is_supported_realtime_input_format", + "json_safe_realtime_payload", + "normalize_conversation_item", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_max_output_tokens", + "realtime_output_format", + "realtime_overlap_fields", + "text_chars_for_audio_ms_from_marks", + "truncate_realtime_item_content", + "validate_conversation_item_audio_formats", + "validate_realtime_item_truncate", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", + "validate_realtime_video_frames", +) + + +@pytest.mark.parametrize("name", _REEXPORTED_FROM_PROTOCOL) +def test_the_duplex_codec_name_is_the_protocol_object(name: str) -> None: + assert getattr(duplex_codec, name) is getattr(protocol, name) + + +@pytest.mark.parametrize( + "name", + ("convert_input_audio_with_rate", "convert_output_audio", "resample_pcm16_mono", "wav_payload_to_pcm16"), +) +def test_the_duplex_audio_shim_is_the_protocol_object(name: str) -> None: + assert getattr(duplex_audio, name) is getattr(protocol, name) + + +def test_a_duplex_command_error_is_a_realtime_protocol_error() -> None: + # The duplex binding converts the codec's error into its own type; both + # carry the same code/event_id, which is what the error envelope renders. + error = DuplexCommandError("bad", code="bad_audio", event_id="event_1") + + assert isinstance(error, RealtimeProtocolError) + assert (error.code, error.event_id) == ("bad_audio", "event_1") + + +def test_the_duplex_append_command_is_built_from_the_shared_decoder() -> None: + defaults = protocol.RealtimeInputDefaults() + event = {"event_id": "event_1", "audio": "", "format": "pcm16", "duration_ms": 40} + + decoded = protocol.decode_audio_append(event, defaults=defaults) + command = duplex_codec.build_append_audio(event, defaults=defaults) + + assert (command.audio, command.format, command.sample_rate_hz) == ( + decoded.audio, + decoded.format, + decoded.sample_rate_hz, + ) + assert (command.is_speech, command.video_frames, command.duration_ms) == ( + decoded.is_speech, + decoded.video_frames, + decoded.duration_ms, + ) + assert (command.audio_end_ms, dict(command.hints), command.event_id) == ( + decoded.audio_end_ms, + decoded.hints, + decoded.event_id, + ) + + +def test_the_shared_decoder_error_reaches_the_client_as_a_duplex_command_error() -> None: + defaults = protocol.RealtimeInputDefaults() + event = {"event_id": "event_1", "audio": "AAAA", "format": "opus"} + + with pytest.raises(RealtimeProtocolError) as shared: + protocol.decode_audio_append(event, defaults=defaults) + with pytest.raises(DuplexCommandError) as duplex: + duplex_codec.build_append_audio(event, defaults=defaults) + + assert (duplex.value.code, str(duplex.value)) == (shared.value.code, str(shared.value)) + assert duplex.value.event_id == shared.value.event_id == "event_1" + + +def test_the_duplex_capabilities_reject_unimplemented_turn_detection() -> None: + # The duplex answer comes from engine.duplex.turn_detection, reached through + # the capability object rather than imported by the codec. + rejection = protocol.validate_session_payload( + {"turn_detection": {"type": "semantic_vad"}}, + capabilities=duplex_codec.DUPLEX_REALTIME_CAPABILITIES, + ) + + assert rejection is not None + assert rejection.code == "unsupported_turn_detection" + + +# ---- the class split (RFC #6592 P0a, second half) ---- +# +# The event and command classes moved out of the engine too. The same identity +# rule applies: the duplex names must BE the protocol objects, and the mailbox +# half must have stayed behind. + +# `docs/serving/realtime_duplex_api.md` sorts every message into three tiers. +# Tier 1 is pure OpenAI and lives in protocol/realtime; Tier 2 (OpenAI names +# carrying our extensions) and Tier 3 (ours alone) live in protocol/duplex. +_TIER1_EVENTS = ("SessionUpdated", "ItemAdded", "ItemCreated", "TextDelta", "OutputItemAdded") +_TIER2_EVENTS = ( + "SessionCreated", + "ResponseCreated", + "ResponseDone", + "AudioDelta", + "InputCommitted", + "ItemDeleted", + "ItemTruncated", + "ErrorEvent", +) +_TIER3_EVENTS = ("Listen", "Speak", "OverlapDecision", "PlaybackAcknowledged", "SessionResumed") +_TIER1_COMMANDS = ("UpdateSession", "CreateItem", "CancelResponse", "ClearInput") +_TIER2_COMMANDS = ("AppendAudio", "Commit") +_TIER3_COMMANDS = ("BargeIn", "AckPlayback", "Heartbeat", "SignalTurn", "CloseSession") + + +@pytest.mark.parametrize("name", _TIER1_EVENTS) +def test_tier1_events_are_the_pure_openai_objects(name: str) -> None: + from vllm_omni.protocol.realtime import events as realtime_events + + assert duplex_events.__dict__[name] is getattr(realtime_events, name) + + +@pytest.mark.parametrize("name", _TIER2_EVENTS + _TIER3_EVENTS) +def test_tier2_and_tier3_events_come_from_protocol_duplex(name: str) -> None: + from vllm_omni.protocol.duplex import events as duplex_wire_events + + assert duplex_events.__dict__[name] is getattr(duplex_wire_events, name) + + +@pytest.mark.parametrize("name", _TIER2_EVENTS) +def test_tier2_extends_its_tier1_twin_rather_than_replacing_it(name: str) -> None: + """A Tier 2 class must be its Tier 1 class plus fields, never a fork of it.""" + from vllm_omni.protocol.duplex import events as duplex_wire_events + from vllm_omni.protocol.realtime import events as realtime_events + + tier1 = getattr(realtime_events, name) + tier2 = getattr(duplex_wire_events, name) + + assert issubclass(tier2, tier1) + assert tier2.wire_type == tier1.wire_type + tier1_fields = {f.name for f in dataclasses.fields(tier1)} + tier2_fields = {f.name for f in dataclasses.fields(tier2)} + # Purely additive: Tier 2 adds, never drops or renames. + assert tier1_fields < tier2_fields + + +@pytest.mark.parametrize("name", _TIER1_EVENTS + _TIER1_COMMANDS) +def test_tier1_carries_no_vllm_omni_extension_fields(name: str) -> None: + """The point of the split: protocol/realtime must be honestly pure OpenAI.""" + from vllm_omni.protocol.realtime import commands as realtime_commands + from vllm_omni.protocol.realtime import events as realtime_events + + cls = getattr(realtime_events, name, None) or getattr(realtime_commands, name) + fields = {f.name for f in dataclasses.fields(cls)} + # Extension fields catalogued as Tier 2 in the serving doc. + assert not ( + fields + & { + "attachment_generation", + "resume_token", + "details", + "extra", + "is_speech", + "video_frames", + "hints", + "realtime_item_id", + } + ) + + +def test_the_error_code_vocabulary_is_tier3() -> None: + """OpenAI standardises the error classes, not our codes (doc: Tier 3).""" + from vllm_omni.protocol.duplex.errors import REALTIME_ERROR_TYPES_BY_CODE + from vllm_omni.protocol.realtime import errors as realtime_errors + + assert duplex_events.REALTIME_ERROR_TYPES_BY_CODE is REALTIME_ERROR_TYPES_BY_CODE + assert not hasattr(realtime_errors, "REALTIME_ERROR_TYPES_BY_CODE") + + +def test_the_tier1_error_event_reports_only_openai_classes() -> None: + from vllm_omni.protocol.duplex import events as duplex_wire_events + from vllm_omni.protocol.realtime import events as realtime_events + + # Tier 1 knows the envelope shape but not our code vocabulary. + assert realtime_events.ErrorEvent(code="resource_exhausted").error_type == "invalid_request_error" + # Tier 2 resolves it through our table. + assert duplex_wire_events.ErrorEvent(code="resource_exhausted").error_type == "rate_limit_error" + + +@pytest.mark.parametrize("name", _TIER1_COMMANDS + _TIER2_COMMANDS + _TIER3_COMMANDS) +def test_engine_commands_subclass_their_protocol_wire_type(name: str) -> None: + from vllm_omni.engine.duplex import commands as engine_commands + from vllm_omni.protocol.duplex import commands as duplex_wire_commands + from vllm_omni.protocol.realtime import commands as realtime_commands + + engine_cls = getattr(engine_commands, name) + wire_cls = getattr(duplex_wire_commands, name, None) or getattr(realtime_commands, name) + + assert issubclass(engine_cls, wire_cls) + assert engine_cls.wire_type == wire_cls.wire_type + + +def test_the_mailbox_channel_stayed_engine_side() -> None: + """``payload()``'s ``type`` is the runner's channel, not the client event. + + These four are the proof that the two halves are genuinely different, and + so the reason the classes were split rather than relocated wholesale. + """ + from vllm_omni.engine.duplex import commands as engine_commands + from vllm_omni.protocol.realtime import commands as realtime_commands + + for name in ("UpdateSession", "CreateItem", "DeleteItem", "TruncateItem"): + engine_cls = getattr(engine_commands, name) + wire_cls = getattr(realtime_commands, name) + assert engine_cls.type == "turn.signal" + assert wire_cls.wire_type != "turn.signal" + # The wire class carries no mailbox vocabulary at all. + assert "type" not in vars(wire_cls) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 959007c30..5f19123bc 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -79,11 +79,12 @@ CHECK_IMPORTS = { "vllm_omni/entrypoints/openai/video_stream_base.py", "vllm_omni/entrypoints/openai/video_stream_session.py", "vllm_omni/experimental/fullduplex/client.py", - "vllm_omni/engine/duplex/audio.py", "vllm_omni/engine/duplex/commands.py", "vllm_omni/engine/duplex/contracts.py", - "vllm_omni/engine/duplex/events.py", - "vllm_omni/engine/duplex/realtime_commands.py", + "vllm_omni/protocol/realtime/audio.py", + "vllm_omni/protocol/realtime/audio_input.py", + "vllm_omni/protocol/realtime/events.py", + "vllm_omni/protocol/realtime/items.py", "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py", diff --git a/vllm_omni/clients/inline_duplex.py b/vllm_omni/clients/inline_duplex.py index 79c09e014..5af8efc5f 100644 --- a/vllm_omni/clients/inline_duplex.py +++ b/vllm_omni/clients/inline_duplex.py @@ -32,8 +32,8 @@ from typing import TYPE_CHECKING from vllm_omni.clients.duplex import DuplexClientBase, DuplexConnectionError, SessionConfig if TYPE_CHECKING: - from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle + from vllm_omni.protocol.duplex import RealtimeInputDefaults __all__ = ["InlineDuplexClient"] @@ -71,7 +71,7 @@ class InlineDuplexClient(DuplexClientBase): def _input_defaults(self) -> RealtimeInputDefaults: """Audio format / sample-rate / VAD defaults declared by this session's object.""" - from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + from vllm_omni.protocol.duplex import RealtimeInputDefaults session_payload = self.session_info or self.config.to_session_payload(model=self.model) return RealtimeInputDefaults().with_session_payload(session_payload) diff --git a/vllm_omni/engine/duplex/audio.py b/vllm_omni/engine/duplex/audio.py index 5f163826a..2d46e8edb 100644 --- a/vllm_omni/engine/duplex/audio.py +++ b/vllm_omni/engine/duplex/audio.py @@ -1,218 +1,44 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compatibility re-export: the Realtime audio codec moved out of the duplex engine. + +Decoding PCM / G.711 / WAV, resampling and re-encoding are wire-format work with +no duplex session in them, so the implementation now lives in +``vllm_omni.protocol.realtime.audio`` where a non-duplex Realtime surface can +use it too. Duplex code reaches it through ``vllm_omni.protocol.duplex``, the +one door the engine uses --- which is also where a duplex-specific conversion +would be introduced. Import from there in new code. +""" + from __future__ import annotations -import base64 -import binascii -import io -import math -import wave +from vllm_omni.protocol.duplex import ( + MAX_INPUT_SAMPLE_RATE_HZ, + MIN_INPUT_SAMPLE_RATE_HZ, + convert_input_audio_with_rate, + convert_output_audio, + decode_g711_alaw, + decode_g711_ulaw, + encode_float32_mono_wav_base64, + encode_g711_alaw, + encode_g711_ulaw, + resample_pcm16_mono, + validate_input_sample_rate_hz, + wav_payload_to_pcm16, +) -import numpy as np - -from vllm_omni.utils.audio_resample import StreamingAudioResampler - -try: - from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin -except ImportError: # pragma: no cover - audioop is removed in newer Python. - alaw2lin = lin2alaw = lin2ulaw = ulaw2lin = None - - -MIN_INPUT_SAMPLE_RATE_HZ = 8_000 -MAX_INPUT_SAMPLE_RATE_HZ = 192_000 - - -def validate_input_sample_rate_hz(sample_rate_hz: object) -> int: - if ( - isinstance(sample_rate_hz, bool) - or not isinstance(sample_rate_hz, int | float) - or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz)) - ): - raise ValueError("sample_rate_hz must be a finite integer") - rate = int(sample_rate_hz) - if rate != sample_rate_hz: - raise ValueError("sample_rate_hz must be an integer") - if not MIN_INPUT_SAMPLE_RATE_HZ <= rate <= MAX_INPUT_SAMPLE_RATE_HZ: - raise ValueError(f"sample_rate_hz must be between {MIN_INPUT_SAMPLE_RATE_HZ} and {MAX_INPUT_SAMPLE_RATE_HZ}") - return rate - - -def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes: - if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz: - return raw - samples = np.frombuffer(raw, dtype=" StreamingAudioResampler.max_polyphase_factor: - # The FIR size scales with the reduced ratio. Keep unusual client - # rates bounded while common audio rates use the high-quality path. - if samples.size <= 1: - return raw - target_size = max(1, int(round(samples.size * target_rate_hz / source_rate_hz))) - source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True) - target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True) - resampled = np.interp(target_x, source_x, samples) - return np.clip(resampled, -32_768, 32_767).astype(" bytes: - if ulaw2lin is not None: - return ulaw2lin(raw, 2) - data = np.frombuffer(raw, dtype=np.uint8) - value = np.bitwise_not(data).astype(np.int16) - sign = value & 0x80 - exponent = (value >> 4) & 0x07 - mantissa = value & 0x0F - sample = ((mantissa << 3) + 0x84) << exponent - return np.where(sign != 0, -(sample - 0x84), sample - 0x84).astype(" bytes: - if alaw2lin is not None: - return alaw2lin(raw, 2) - data = np.bitwise_xor(np.frombuffer(raw, dtype=np.uint8), 0x55).astype(np.int16) - sign = data & 0x80 - exponent = (data >> 4) & 0x07 - mantissa = data & 0x0F - sample = np.where(exponent == 0, (mantissa << 4) + 8, ((mantissa << 4) + 0x108) << (exponent - 1)) - return np.where(sign != 0, sample, -sample).astype(" bytes: - if lin2ulaw is not None: - return lin2ulaw(raw, 2) - pcm = np.clip(np.frombuffer(raw, dtype=" (0xFF << exp), exp + 1, exponent) - mantissa = (magnitude >> (exponent + 3)) & 0x0F - return (np.bitwise_not(sign | (exponent << 4) | mantissa) & 0xFF).astype(np.uint8).tobytes() - - -def encode_g711_alaw(raw: bytes) -> bytes: - if lin2alaw is not None: - return lin2alaw(raw, 2) - pcm = np.frombuffer(raw, dtype="= 0, 0x80, 0x00) - magnitude = np.abs(pcm) - exponent = np.zeros_like(magnitude) - for exp in range(1, 8): - exponent = np.where(magnitude >= (1 << (exp + 7)), exp, exponent) - mantissa = np.where( - exponent == 0, - (magnitude >> 4) & 0x0F, - (magnitude >> (exponent + 3)) & 0x0F, - ) - return ((sign | (exponent << 4) | mantissa) ^ 0x55).astype(np.uint8).tobytes() - - -def wav_payload_to_pcm16(raw: bytes) -> tuple[bytes | None, int | None]: - try: - with wave.open(io.BytesIO(raw), "rb") as wav_file: - channels = wav_file.getnchannels() - sample_width = wav_file.getsampwidth() - sample_rate_hz = wav_file.getframerate() - frames = wav_file.readframes(wav_file.getnframes()) - if sample_width != 2: - return None, sample_rate_hz - if channels <= 1: - return frames, sample_rate_hz - pcm = np.frombuffer(frames, dtype=" str: - """Package normalized mono float32 samples as a PCM16 WAV data payload.""" - rate = validate_input_sample_rate_hz(sample_rate_hz) - normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) - pcm16 = np.clip(np.rint(normalized * 32768.0), -32768, 32767).astype(" tuple[object, object, int | float | None]: - if not isinstance(audio, str) or not isinstance(fmt, str): - return audio, fmt, sample_rate_hz - normalized = fmt.lower() - if normalized not in {"pcm16", "pcm_s16le", "s16le", "g711_ulaw", "g711_alaw"}: - return audio, fmt, sample_rate_hz - if isinstance(sample_rate_hz, int | float): - sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) - try: - raw = base64.b64decode(audio.strip(), validate=False) - except (binascii.Error, ValueError): - return audio, fmt, sample_rate_hz - if normalized == "g711_ulaw": - raw = decode_g711_ulaw(raw) - sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 - elif normalized == "g711_alaw": - raw = decode_g711_alaw(raw) - sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 - elif len(raw) % 2: - return audio, fmt, sample_rate_hz - if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) != target_sample_rate_hz: - raw = resample_pcm16_mono( - raw, - source_rate_hz=int(sample_rate_hz), - target_rate_hz=target_sample_rate_hz, - ) - sample_rate_hz = target_sample_rate_hz - pcm = np.frombuffer(raw, dtype=" tuple[str, str, int | None]: - target = target_fmt.lower() - if target not in {"g711_ulaw", "g711_alaw"}: - return audio, source_fmt, source_sample_rate_hz - try: - raw = base64.b64decode(audio, validate=False) - except (binascii.Error, ValueError): - return audio, source_fmt, source_sample_rate_hz - source = source_fmt.lower() - if source == "wav": - pcm_raw, wav_rate = wav_payload_to_pcm16(raw) - if pcm_raw is None: - return audio, source_fmt, source_sample_rate_hz - raw = pcm_raw - source_sample_rate_hz = source_sample_rate_hz or wav_rate - elif source not in {"pcm", "pcm16", "pcm_s16le", "s16le"} or len(raw) % 2: - return audio, source_fmt, source_sample_rate_hz - target_rate = target_sample_rate_hz or 8_000 - if source_sample_rate_hz is not None: - raw = resample_pcm16_mono(raw, source_rate_hz=source_sample_rate_hz, target_rate_hz=target_rate) - encoded = encode_g711_ulaw(raw) if target == "g711_ulaw" else encode_g711_alaw(raw) - return base64.b64encode(encoded).decode("ascii"), target, target_rate +__all__ = [ + "MAX_INPUT_SAMPLE_RATE_HZ", + "MIN_INPUT_SAMPLE_RATE_HZ", + "convert_input_audio_with_rate", + "convert_output_audio", + "decode_g711_alaw", + "decode_g711_ulaw", + "encode_float32_mono_wav_base64", + "encode_g711_alaw", + "encode_g711_ulaw", + "resample_pcm16_mono", + "validate_input_sample_rate_hz", + "wav_payload_to_pcm16", +] diff --git a/vllm_omni/engine/duplex/commands.py b/vllm_omni/engine/duplex/commands.py index 90d48d050..19674948c 100644 --- a/vllm_omni/engine/duplex/commands.py +++ b/vllm_omni/engine/duplex/commands.py @@ -1,41 +1,62 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Typed commands accepted by a duplex session. +"""The command vocabulary one duplex session accepts, plus its mailbox rendering. ``DuplexSessionHandle.submit()`` takes one of these; the websocket handler and -``InlineDuplexClient`` build them with :func:`command_from_realtime` from the -OpenAI Realtime client event vocabulary. Every command can also render the -session-internal mailbox payload (``payload()``), which is the dictionary -vocabulary the session runner bodies were written against. +``InlineDuplexClient`` build them with :func:`command_from_realtime`. + +The *wire* half of each command --- which client event it decodes from and what +fields survive decoding --- now lives in ``vllm_omni.protocol.duplex.commands``, +which carries the whole vocabulary: the eight Tier 1 classes re-exported from +``vllm_omni.protocol.realtime.commands``, the two Tier 2 ones and the seven +Tier 3 ones. The engine imports that module and never the Tier 1 one directly, +so a command that later grows a duplex extension changes one file +(RFC #6592 P0a). + +The *engine* half stays here, because it is not wire contract at all: +``payload()`` renders the session-internal mailbox dictionary the runner bodies +were written against, and its ``type`` is the mailbox channel rather than the +client event. Those two genuinely differ --- ``session.update`` and +``conversation.item.create`` / ``.delete`` / ``.truncate`` all travel on the +``turn.signal`` channel --- which is why the halves are separated instead of +the classes being relocated wholesale. """ from __future__ import annotations import base64 from collections.abc import Mapping -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, ClassVar +from vllm_omni.protocol.duplex import RealtimeProtocolError +from vllm_omni.protocol.duplex import commands as _duplex_wire +from vllm_omni.protocol.duplex.commands import RealtimeCommand + if TYPE_CHECKING: - from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + from vllm_omni.protocol.duplex import RealtimeInputDefaults -class DuplexCommandError(ValueError): - """A client payload could not be turned into a command.""" +class DuplexCommandError(RealtimeProtocolError): + """A client payload could not be turned into a duplex command. - def __init__(self, message: str, *, code: str = "bad_event", event_id: str | None = None) -> None: - super().__init__(message) - self.code = code - self.event_id = event_id + The duplex name for a Realtime protocol error: same ``code`` / + ``event_id`` contract, so the error envelope is rendered identically + whichever consumer raised it. + """ @dataclass(frozen=True, slots=True, kw_only=True) -class DuplexCommand: +class DuplexCommand(RealtimeCommand): + """A Realtime command as the duplex engine handles it. + + Adds the mailbox channel (``type``) and its rendering on top of the wire + command; every concrete class below pairs this with its protocol twin. + """ + #: Mailbox event type this command renders to (see ``payload()``). type: ClassVar[str] = "" - #: Client correlation id (OpenAI ``event_id``), echoed on error events. - event_id: str | None = None def payload(self) -> dict[str, object]: """Render the session-internal mailbox dictionary.""" @@ -56,19 +77,23 @@ class DuplexCommand: return data +# ---- GA commands ---- + + @dataclass(frozen=True, slots=True, kw_only=True) -class AppendAudio(DuplexCommand): +class UpdateSession(DuplexCommand, _duplex_wire.UpdateSession): + type: ClassVar[str] = "turn.signal" + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "session.update" + data["payload"] = dict(data.pop("patch", {}) or {}) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendAudio(DuplexCommand, _duplex_wire.AppendAudio): type: ClassVar[str] = "input_audio_buffer.append" - #: Raw audio bytes in ``format`` at ``sample_rate_hz`` (base64 only on the wire). - audio: bytes - format: str = "pcm16" - sample_rate_hz: int | None = None - is_speech: bool | None = None - video_frames: tuple[str, ...] = () - duration_ms: int | None = None - audio_end_ms: int | None = None - #: Model-neutral hints carried through from the wire (rms, vad, transcript hints ...). - hints: Mapping[str, object] = field(default_factory=dict) def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -90,21 +115,8 @@ class AppendAudio(DuplexCommand): @dataclass(frozen=True, slots=True, kw_only=True) -class AppendText(DuplexCommand): - type: ClassVar[str] = "input.text.append" - text: str - - -@dataclass(frozen=True, slots=True, kw_only=True) -class Commit(DuplexCommand): +class Commit(DuplexCommand, _duplex_wire.Commit): type: ClassVar[str] = "input_audio_buffer.commit" - final: bool = True - #: ``None`` means "no explicit request": the runner decides on commit (auto-response - #: sessions answer on their own); ``True`` / ``False`` force it. - create_response: bool | None = None - is_speech: bool | None = None - #: Realtime conversation item created for this commit (wire correlation only). - realtime_item_id: str | None = None def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -114,10 +126,18 @@ class Commit(DuplexCommand): @dataclass(frozen=True, slots=True, kw_only=True) -class CreateResponse(DuplexCommand): +class ClearInput(DuplexCommand, _duplex_wire.ClearInput): + type: ClassVar[str] = "input_audio_buffer.clear" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearOutputAudio(DuplexCommand, _duplex_wire.ClearOutputAudio): + type: ClassVar[str] = "output_audio_buffer.clear" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateResponse(DuplexCommand, _duplex_wire.CreateResponse): type: ClassVar[str] = "response.create" - #: Raw Realtime ``response`` object (instructions, voice, modalities, ...). - options: Mapping[str, object] = field(default_factory=dict) def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -128,84 +148,13 @@ class CreateResponse(DuplexCommand): @dataclass(frozen=True, slots=True, kw_only=True) -class ClearInput(DuplexCommand): - type: ClassVar[str] = "input_audio_buffer.clear" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class CancelInput(DuplexCommand): - type: ClassVar[str] = "input.cancel" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class CancelResponse(DuplexCommand): +class CancelResponse(DuplexCommand, _duplex_wire.CancelResponse): type: ClassVar[str] = "response.cancel" - response_id: str | None = None @dataclass(frozen=True, slots=True, kw_only=True) -class BargeIn(DuplexCommand): - type: ClassVar[str] = "barge_in" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ClearOutputAudio(DuplexCommand): - type: ClassVar[str] = "output_audio_buffer.clear" - #: Explicit response to clear; ``None`` targets the active/last response. - response_id: str | None = None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SignalTurn(DuplexCommand): - """Generic ``turn.signal`` (local turn transitions such as ``user_started``).""" - +class CreateItem(DuplexCommand, _duplex_wire.CreateItem): type: ClassVar[str] = "turn.signal" - event: str - signal_payload: Mapping[str, object] = field(default_factory=dict) - - def payload(self) -> dict[str, object]: - data = DuplexCommand.payload(self) - signal_payload = data.pop("signal_payload", None) - if isinstance(signal_payload, Mapping) and signal_payload: - data["payload"] = dict(signal_payload) - return data - - -@dataclass(frozen=True, slots=True, kw_only=True) -class UpdateSession(DuplexCommand): - """``session.update``: a Realtime ``session`` object patch.""" - - type: ClassVar[str] = "turn.signal" - patch: Mapping[str, object] = field(default_factory=dict) - - def payload(self) -> dict[str, object]: - data = DuplexCommand.payload(self) - data["event"] = "session.update" - data["payload"] = dict(data.pop("patch", {}) or {}) - return data - - -@dataclass(frozen=True, slots=True, kw_only=True) -class AckPlayback(DuplexCommand): - type: ClassVar[str] = "playback.ack" - played_ms: int - committed_ms: int | None = None - response_id: str | None = None - item_id: str | None = None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class Heartbeat(DuplexCommand): - type: ClassVar[str] = "session.heartbeat" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class CreateItem(DuplexCommand): - """``conversation.item.create`` (history injection or function-call output).""" - - type: ClassVar[str] = "turn.signal" - item: Mapping[str, object] - previous_item_id: str | None = None def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -219,9 +168,8 @@ class CreateItem(DuplexCommand): @dataclass(frozen=True, slots=True, kw_only=True) -class DeleteItem(DuplexCommand): +class DeleteItem(DuplexCommand, _duplex_wire.DeleteItem): type: ClassVar[str] = "turn.signal" - item_id: str def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -231,11 +179,8 @@ class DeleteItem(DuplexCommand): @dataclass(frozen=True, slots=True, kw_only=True) -class TruncateItem(DuplexCommand): +class TruncateItem(DuplexCommand, _duplex_wire.TruncateItem): type: ClassVar[str] = "turn.signal" - item_id: str - audio_end_ms: int - content_index: int = 0 def payload(self) -> dict[str, object]: data = DuplexCommand.payload(self) @@ -248,15 +193,51 @@ class TruncateItem(DuplexCommand): return data +# ---- duplex extension commands ---- + + @dataclass(frozen=True, slots=True, kw_only=True) -class CloseSession(DuplexCommand): - """Graceful close requested through the command stream (``session.close``).""" +class AppendText(DuplexCommand, _duplex_wire.AppendText): + type: ClassVar[str] = "input.text.append" + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelInput(DuplexCommand, _duplex_wire.CancelInput): + type: ClassVar[str] = "input.cancel" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BargeIn(DuplexCommand, _duplex_wire.BargeIn): + type: ClassVar[str] = "barge_in" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SignalTurn(DuplexCommand, _duplex_wire.SignalTurn): + type: ClassVar[str] = "turn.signal" + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + signal_payload = data.pop("signal_payload", None) + if isinstance(signal_payload, Mapping) and signal_payload: + data["payload"] = dict(signal_payload) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AckPlayback(DuplexCommand, _duplex_wire.AckPlayback): + type: ClassVar[str] = "playback.ack" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Heartbeat(DuplexCommand, _duplex_wire.Heartbeat): + type: ClassVar[str] = "session.heartbeat" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CloseSession(DuplexCommand, _duplex_wire.CloseSession): type: ClassVar[str] = "session.close" - reason: str = "client_close" -#: Realtime client event types that map onto commands (wire vocabulary). REALTIME_COMMAND_TYPES: frozenset[str] = frozenset( { "input_audio_buffer.append", @@ -300,7 +281,6 @@ def command_from_realtime( __all__ = [ - "REALTIME_COMMAND_TYPES", "AckPlayback", "AppendAudio", "AppendText", @@ -317,6 +297,8 @@ __all__ = [ "DuplexCommand", "DuplexCommandError", "Heartbeat", + "REALTIME_COMMAND_TYPES", + "RealtimeCommand", "SignalTurn", "TruncateItem", "UpdateSession", diff --git a/vllm_omni/engine/duplex/config.py b/vllm_omni/engine/duplex/config.py index aab0fb1c5..a898a2cf5 100644 --- a/vllm_omni/engine/duplex/config.py +++ b/vllm_omni/engine/duplex/config.py @@ -339,26 +339,26 @@ class DuplexSessionConfig: ``_session_create_from_realtime``. """ from vllm_omni.engine.duplex.realtime_commands import ( - RealtimeInputDefaults, + DUPLEX_REALTIME_CAPABILITIES, duplex_response_format, + ) + from vllm_omni.engine.duplex.turn_detection import normalize_turn_detection_session_payload + from vllm_omni.protocol.duplex import ( + RealtimeInputDefaults, input_audio_transcription_config, json_safe_realtime_payload, realtime_max_output_tokens, realtime_overlap_fields, - validate_realtime_session_audio_formats, - ) - from vllm_omni.engine.duplex.turn_detection import ( - normalize_turn_detection_session_payload, - validate_realtime_turn_detection, + validate_session_payload, ) payload: dict[str, object] = dict(session_payload) - format_error = validate_realtime_session_audio_formats(payload) - if format_error is not None: - raise DuplexConfigError(format_error, code="unsupported_audio_format") - turn_detection_error = validate_realtime_turn_detection(payload) - if turn_detection_error is not None: - raise DuplexConfigError(turn_detection_error, code="unsupported_turn_detection", param="turn_detection") + # One session check for every consumer (ENTRY-INV-002): the same + # capability object ``translate_realtime_command`` uses, so a session + # object is accepted or refused identically whichever door it came in. + rejection = validate_session_payload(payload, capabilities=DUPLEX_REALTIME_CAPABILITIES) + if rejection is not None: + raise DuplexConfigError(rejection.message, code=rejection.code, param=rejection.param) normalize_turn_detection_session_payload(payload) defaults = RealtimeInputDefaults().with_session_payload(payload) payload.update(realtime_overlap_fields(payload)) @@ -446,9 +446,9 @@ class DuplexSessionConfig: when the patch changes something a live session cannot change. ``audio_started`` is ``playback.generated_ms > 0 or playback.sent_ms > 0``. """ - from vllm_omni.engine.duplex.realtime_commands import ( + from vllm_omni.engine.duplex.realtime_commands import duplex_response_format + from vllm_omni.protocol.duplex import ( REALTIME_OUTPUT_AUDIO_FORMATS, - duplex_response_format, input_audio_transcription_config, json_safe_realtime_payload, parse_realtime_audio_format, @@ -629,9 +629,9 @@ class ResponseCreateOptions: for options a model-native duplex session cannot apply per response. Private runtime keys in ``extra_body`` are dropped. """ - from vllm_omni.engine.duplex.realtime_commands import ( + from vllm_omni.engine.duplex.realtime_commands import duplex_response_format + from vllm_omni.protocol.duplex import ( REALTIME_OUTPUT_AUDIO_FORMATS, - duplex_response_format, parse_realtime_audio_format, realtime_max_output_tokens, ) @@ -782,14 +782,14 @@ def realtime_item_to_history_message(item: object) -> dict[str, object] | None: def realtime_max_output_tokens(value: object) -> int | None: """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" - from vllm_omni.engine.duplex.realtime_commands import realtime_max_output_tokens as _impl + from vllm_omni.protocol.duplex import realtime_max_output_tokens as _impl return _impl(value) def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: """Return the ``input_audio_transcription`` object of a Realtime session payload.""" - from vllm_omni.engine.duplex.realtime_commands import input_audio_transcription_config as _impl + from vllm_omni.protocol.duplex import input_audio_transcription_config as _impl return _impl(session_payload) diff --git a/vllm_omni/engine/duplex/events.py b/vllm_omni/engine/duplex/events.py index 9c131a893..197dec613 100644 --- a/vllm_omni/engine/duplex/events.py +++ b/vllm_omni/engine/duplex/events.py @@ -1,677 +1,79 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Typed events emitted by a duplex session. +"""The event vocabulary one duplex session speaks, assembled from two layers. -Every public event is a frozen dataclass with explicit fields. The Realtime -wire JSON is *derived* from those fields by :meth:`DuplexEvent.to_realtime`, -which is pure: it never consults session state. The stateful part (response / -item ids, content-part bookkeeping) lives in -``vllm_omni.engine.duplex.realtime_events`` on the session runner and is -consumed when the events are *constructed*. +The classes themselves no longer live here. ``vllm_omni.protocol.duplex.events`` +carries the whole vocabulary --- 22 Tier 1 classes re-exported from +``vllm_omni.protocol.realtime.events``, 8 Tier 2 ones and 12 Tier 3 ones --- so +a non-duplex Realtime consumer can take the Tier 1 half and its wire rendering +without the duplex session control plane (RFC #6592 P0a). The engine imports +the duplex module and never the Tier 1 one directly. -``session_id`` / ``epoch`` are bound by -``DuplexSessionManager.emit`` when an event leaves the runner; producers -(projector, serving handshake helpers) build events with the defaults. +What is genuinely engine-side stays: :data:`DOMAIN_TERMINAL_EVENTS` and +:data:`MODEL_OUTPUT_EVENTS` are the session runner's epoch-filter policy, not +wire contract --- they name *internal* event types, and which of them may never +be dropped as stale is a runtime decision. + +``DuplexEvent`` is re-exported as the name the engine and the clients already +use. It is an alias of +:class:`~vllm_omni.protocol.realtime.events.RealtimeEvent`, not a subclass, +because an event has no engine-internal half --- what the session emits is +exactly what the socket sends. A *command* does have one (``payload()`` renders +the runner's mailbox), which is why ``DuplexCommand`` is a real class. """ from __future__ import annotations -import base64 -from collections.abc import Mapping -from dataclasses import dataclass, field, fields -from typing import ClassVar, cast -from uuid import uuid4 +from vllm_omni.protocol.duplex.errors import REALTIME_ERROR_TYPES_BY_CODE +from vllm_omni.protocol.duplex.events import ( + AudioDelta, + AudioDone, + ContentPartAdded, + ContentPartDone, + DuplexEvent, + DuplexRawEvent, + ErrorEvent, + FunctionCallArgumentsDelta, + FunctionCallArgumentsDone, + InputCleared, + InputCommitted, + InputTranscriptionCompleted, + ItemAdded, + ItemCreated, + ItemDeleted, + ItemDone, + ItemRetrieved, + ItemTruncated, + Listen, + OutputAudioCleared, + OutputItemAdded, + OutputItemDone, + OverlapDecision, + PlaybackAcknowledged, + RateLimitsUpdated, + ResponseCreated, + ResponseDone, + SessionClosed, + SessionCreated, + SessionExpired, + SessionHeartbeatAck, + SessionReplaced, + SessionResumed, + SessionResyncRequired, + SessionUpdated, + Speak, + SpeechStarted, + SpeechStopped, + TextDelta, + TextDone, + TranscriptDelta, + TranscriptDone, + TurnEvent, + error_event, + new_event_id, +) - -def new_event_id() -> str: - return f"event_{uuid4().hex}" - - -#: OpenAI Realtime ``error.type`` for each internal error code. -REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = { - "bad_event": "invalid_request_error", - "bad_audio": "invalid_request_error", - "config_timeout": "invalid_request_error", - "invalid_json": "invalid_request_error", - "event_too_large": "invalid_request_error", - "unknown_event": "invalid_request_error", - "internal_error": "server_error", - "runtime_append_failed": "server_error", - "runtime_append_task_failed": "server_error", - "runtime_signal_failed": "server_error", - "runtime_abort_failed": "server_error", - "runtime_data_plane_stream_failed": "server_error", - "runtime_data_plane_text_without_audio": "server_error", - "resource_exhausted": "rate_limit_error", - "session_exists": "invalid_request_error", - "session_closed": "invalid_request_error", - "unknown_session": "invalid_request_error", - "invalid_duplex_runtime_config": "invalid_request_error", - "instructions_update_unsupported": "invalid_request_error", - "persona_update_unsupported": "invalid_request_error", - "voice_update_unsupported": "invalid_request_error", - "unsupported_nemotron_duplex_mode": "invalid_request_error", - "unsupported_native_response_options": "invalid_request_error", - "runtime_touch_failed": "server_error", - "engine_error": "server_error", - "input_backpressure": "rate_limit_error", - "response_already_active": "invalid_request_error", - "response_not_active": "invalid_request_error", - "response_create_without_input": "invalid_request_error", - "text_only_turn_unsupported": "invalid_request_error", - "commit_aborted": "server_error", - "input_audio_buffer_empty": "invalid_request_error", - "missing_item_id": "invalid_request_error", - "item_not_found": "invalid_request_error", - "playback_item_mismatch": "invalid_request_error", - "playback_item_not_found": "invalid_request_error", - "playback_ack_too_late": "invalid_request_error", - "unsupported_audio_format": "invalid_request_error", - "unsupported_turn_detection": "invalid_request_error", - "unsupported_ref_audio_path": "invalid_request_error", - "ref_audio_required": "invalid_request_error", - "model_update_unsupported": "invalid_request_error", - "voice_update_after_audio_unsupported": "invalid_request_error", - "ref_audio_update_unsupported": "invalid_request_error", - "native_text_append_unsupported": "invalid_request_error", - "invalid_video_frames": "invalid_request_error", - "invalid_function_call_output": "invalid_request_error", - "server_vad_unavailable": "server_error", -} - -_IDENTITY_FIELDS = frozenset({"session_id", "epoch", "event_id"}) - - -def _wire_value(value: object) -> object: - if isinstance(value, Mapping): - return {str(k): _wire_value(v) for k, v in value.items()} - if isinstance(value, tuple | list): - return [_wire_value(v) for v in value] - return value - - -@dataclass(frozen=True, slots=True, kw_only=True) -class DuplexEvent: - """Base of every public session event.""" - - #: Realtime ``type`` this class renders to. - wire_type: ClassVar[str] = "" - #: Fields that are omitted from the wire object when ``None``. - optional_wire_fields: ClassVar[frozenset[str]] = frozenset() - - session_id: str = "" - epoch: int | None = None - #: Server event id (OpenAI ``event_id``); generated at construction. - event_id: str = field(default_factory=new_event_id) - - # ---- identity ---- - - @property - def type(self) -> str: - return self.wire_type - - @property - def is_terminal(self) -> bool: - return False - - # ---- generic accessors (None when the class has no such field) ---- - - @property - def response_id(self) -> str | None: - return None - - @property - def item_id(self) -> str | None: - return None - - @property - def text(self) -> str | None: - return None - - @property - def audio(self) -> bytes | None: - return None - - # ---- wire rendering (pure) ---- - - def _wire_fields(self) -> dict[str, object]: - data: dict[str, object] = {} - for f in fields(self): - if f.name in _IDENTITY_FIELDS: - continue - value = getattr(self, f.name) - if value is None and f.name in self.optional_wire_fields: - continue - data[f.name] = _wire_value(value) - return data - - def to_realtime(self) -> dict[str, object]: - """The Realtime wire JSON object for this event (derived from the fields).""" - return {"type": self.wire_type, "event_id": self.event_id, **self._wire_fields()} - - def __repr__(self) -> str: # pragma: no cover - debugging aid - return f"{type(self).__name__}(type={self.wire_type!r}, session_id={self.session_id!r}, epoch={self.epoch!r})" - - -# ---- session lifecycle ---- - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionCreated(DuplexEvent): - wire_type = "session.created" - optional_wire_fields = frozenset({"attachment_generation", "resume_token"}) - - session: Mapping[str, object] = field(default_factory=dict) - #: Transport resume credentials (set by the websocket handler when resume is supported). - attachment_generation: int | None = None - resume_token: str | None = None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionUpdated(DuplexEvent): - wire_type = "session.updated" - - session: Mapping[str, object] = field(default_factory=dict) - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionResumed(SessionCreated): - wire_type = "session.resumed" - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionReplaced(DuplexEvent): - """A newer connection took the session over (sent to the replaced socket).""" - - wire_type = "session.replaced" - - attachment_generation: int = 0 - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionResyncRequired(DuplexEvent): - """The replay journal cannot bridge the gap; the client must start over.""" - - wire_type = "session.resync_required" - - reason: str = "journal_gap" - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionClosed(DuplexEvent): - wire_type = "session.closed" - - reason: str = "closed" - #: The session-internal close event (kept for clients that inspect it). - details: Mapping[str, object] = field(default_factory=dict) - - @property - def is_terminal(self) -> bool: - return True - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, "reason": self.reason, "event": _wire_value(self.details)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionExpired(SessionClosed): - wire_type = "session.expired" - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, "reason": self.reason} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SessionHeartbeatAck(DuplexEvent): - wire_type = "session.heartbeat_ack" - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class TurnEvent(DuplexEvent): - """Turn-state transition (``turn.event``).""" - - wire_type = "turn.event" - - event: str = "" - turn_state: str = "" - - def _wire_fields(self) -> dict[str, object]: - return {"session_id": self.session_id, "event": self.event, "turn_state": self.turn_state, "epoch": self.epoch} - - -# ---- response lifecycle ---- - - -@dataclass(frozen=True, slots=True, kw_only=True) -class _ResponseEvent(DuplexEvent): - """Events addressed to one response (and usually one output item / content part).""" - - response_id: str | None = None # type: ignore[assignment] - item_id: str | None = None # type: ignore[assignment] - output_index: int = 0 - content_index: int = 0 - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ResponseCreated(DuplexEvent): - wire_type = "response.created" - - response_id: str | None = None # type: ignore[assignment] - response: Mapping[str, object] = field(default_factory=dict) - - -@dataclass(frozen=True, slots=True, kw_only=True) -class Listen(DuplexEvent): - """The model decided to keep listening (no spoken response for this turn).""" - - wire_type = "response.listen" - optional_wire_fields = frozenset({"response_id"}) - - response_id: str | None = None # type: ignore[assignment] - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - response: dict[str, object] = { - "object": "realtime.response", - "status": "listening", - "metadata": _wire_value(self.details), - } - data: dict[str, object] = {"session_id": self.session_id, "epoch": self.epoch, "response": response} - if self.response_id: - response["id"] = self.response_id - data["response_id"] = self.response_id - return data - - -@dataclass(frozen=True, slots=True, kw_only=True) -class Speak(_ResponseEvent): - wire_type = "response.speak" - - metadata: Mapping[str, object] = field(default_factory=dict) - - -@dataclass(frozen=True, slots=True, kw_only=True) -class OverlapDecision(DuplexEvent): - wire_type = "overlap.decision" - - policy: str | None = None - action: str | None = None - reason: str | None = None - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - return { - "session_id": self.session_id, - "epoch": self.epoch, - "policy": self.policy, - "action": self.action, - "reason": self.reason, - "metadata": _wire_value(self.details), - } - - -@dataclass(frozen=True, slots=True, kw_only=True) -class OutputItemAdded(DuplexEvent): - wire_type = "response.output_item.added" - - response_id: str | None = None # type: ignore[assignment] - output_index: int = 0 - item: Mapping[str, object] = field(default_factory=dict) - - @property - def item_id(self) -> str | None: - value = self.item.get("id") - return value if isinstance(value, str) else None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class OutputItemDone(OutputItemAdded): - wire_type = "response.output_item.done" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ContentPartAdded(_ResponseEvent): - wire_type = "response.content_part.added" - - part: Mapping[str, object] = field(default_factory=dict) - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ContentPartDone(ContentPartAdded): - wire_type = "response.content_part.done" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class AudioDelta(_ResponseEvent): - wire_type = "response.output_audio.delta" - optional_wire_fields = frozenset({"sample_rate_hz", "metadata"}) - - #: Base64 audio in ``format``. - delta: str = "" - format: str = "pcm16" - sample_rate_hz: int | None = None - metadata: Mapping[str, object] | None = None - - @property - def audio(self) -> bytes | None: - try: - return base64.b64decode(self.delta) - except (ValueError, TypeError): - return None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class AudioDone(_ResponseEvent): - wire_type = "response.output_audio.done" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class TranscriptDelta(_ResponseEvent): - wire_type = "response.output_audio_transcript.delta" - - delta: str = "" - - @property - def text(self) -> str | None: - return self.delta - - -@dataclass(frozen=True, slots=True, kw_only=True) -class TranscriptDone(_ResponseEvent): - wire_type = "response.output_audio_transcript.done" - - transcript: str = "" - - @property - def text(self) -> str | None: - return self.transcript - - -@dataclass(frozen=True, slots=True, kw_only=True) -class TextDelta(_ResponseEvent): - wire_type = "response.output_text.delta" - - delta: str = "" - - @property - def text(self) -> str | None: - return self.delta - - -@dataclass(frozen=True, slots=True, kw_only=True) -class TextDone(_ResponseEvent): - wire_type = "response.output_text.done" - - text: str = "" # type: ignore[assignment] - - -@dataclass(frozen=True, slots=True, kw_only=True) -class FunctionCallArgumentsDelta(DuplexEvent): - wire_type = "response.function_call_arguments.delta" - - response_id: str | None = None # type: ignore[assignment] - item_id: str | None = None # type: ignore[assignment] - output_index: int = 0 - call_id: str = "" - delta: str = "" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class FunctionCallArgumentsDone(DuplexEvent): - wire_type = "response.function_call_arguments.done" - - response_id: str | None = None # type: ignore[assignment] - item_id: str | None = None # type: ignore[assignment] - output_index: int = 0 - call_id: str = "" - arguments: str = "" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ResponseDone(DuplexEvent): - wire_type = "response.done" - - response_id: str | None = None # type: ignore[assignment] - response: Mapping[str, object] = field(default_factory=dict) - - @property - def status(self) -> str | None: - value = self.response.get("status") - return value if isinstance(value, str) else None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class RateLimitsUpdated(DuplexEvent): - wire_type = "rate_limits.updated" - - rate_limits: tuple[Mapping[str, object], ...] = () - - -# ---- input buffer / conversation ---- - - -@dataclass(frozen=True, slots=True, kw_only=True) -class InputCommitted(DuplexEvent): - wire_type = "input_audio_buffer.committed" - - previous_item_id: str | None = None - item_id: str | None = None # type: ignore[assignment] - #: The session-internal ``input.committed`` event. - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - return { - "previous_item_id": self.previous_item_id, - "item_id": self.item_id, - "event": _wire_value(self.details), - } - - -@dataclass(frozen=True, slots=True, kw_only=True) -class InputCleared(DuplexEvent): - wire_type = "input_audio_buffer.cleared" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SpeechStarted(DuplexEvent): - wire_type = "input_audio_buffer.speech_started" - - audio_start_ms: int = 0 - item_id: str | None = None # type: ignore[assignment] - - -@dataclass(frozen=True, slots=True, kw_only=True) -class SpeechStopped(DuplexEvent): - wire_type = "input_audio_buffer.speech_stopped" - - audio_end_ms: int = 0 - item_id: str | None = None # type: ignore[assignment] - - -@dataclass(frozen=True, slots=True, kw_only=True) -class _ItemEvent(DuplexEvent): - previous_item_id: str | None = None - item: Mapping[str, object] = field(default_factory=dict) - - @property - def item_id(self) -> str | None: - value = self.item.get("id") - return value if isinstance(value, str) else None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemAdded(_ItemEvent): - wire_type = "conversation.item.added" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemCreated(_ItemEvent): - wire_type = "conversation.item.created" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemDone(_ItemEvent): - wire_type = "conversation.item.done" - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemRetrieved(DuplexEvent): - wire_type = "conversation.item.retrieved" - - item: Mapping[str, object] = field(default_factory=dict) - - @property - def item_id(self) -> str | None: - value = self.item.get("id") - return value if isinstance(value, str) else None - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemDeleted(DuplexEvent): - wire_type = "conversation.item.deleted" - - item_id: str | None = None # type: ignore[assignment] - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - return {"item_id": self.item_id, "event": _wire_value(self.details)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ItemTruncated(DuplexEvent): - wire_type = "conversation.item.truncated" - - item_id: str | None = None # type: ignore[assignment] - content_index: int = 0 - audio_end_ms: int = 0 - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - return { - "item_id": self.item_id, - "content_index": self.content_index, - "audio_end_ms": self.audio_end_ms, - "event": _wire_value(self.details), - } - - -@dataclass(frozen=True, slots=True, kw_only=True) -class InputTranscriptionCompleted(DuplexEvent): - wire_type = "conversation.item.input_audio_transcription.completed" - - item_id: str | None = None # type: ignore[assignment] - content_index: int = 0 - transcript: str = "" - - @property - def text(self) -> str | None: - return self.transcript - - -# ---- playback / control ---- - - -@dataclass(frozen=True, slots=True, kw_only=True) -class OutputAudioCleared(DuplexEvent): - wire_type = "output_audio_buffer.cleared" - - response_id: str | None = None # type: ignore[assignment] - - -@dataclass(frozen=True, slots=True, kw_only=True) -class PlaybackAcknowledged(DuplexEvent): - wire_type = "playback.acknowledged" - - #: The session-internal ``playback.acknowledged`` event (cursor, committed_ms ...). - details: Mapping[str, object] = field(default_factory=dict) - - def _wire_fields(self) -> dict[str, object]: - return {"event": _wire_value(self.details)} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class ErrorEvent(DuplexEvent): - wire_type = "error" - - code: str = "internal_error" - message: str = "" - #: Client ``event_id`` this error answers (OpenAI ``error.event_id``). - related_event_id: str | None = None - param: str | None = None - #: Extra keys merged into the wire ``error`` object (``retryable`` ...). - extra: Mapping[str, object] = field(default_factory=dict) - - @property - def error_type(self) -> str: - return REALTIME_ERROR_TYPES_BY_CODE.get(self.code, "invalid_request_error") - - @property - def error(self) -> dict[str, object]: - error: dict[str, object] = {"type": self.error_type, "code": self.code, "message": self.message} - if self.related_event_id: - error["event_id"] = self.related_event_id - if self.param: - error["param"] = self.param - error.update(cast("Mapping[str, object]", _wire_value(self.extra))) - return error - - @property - def text(self) -> str | None: - return self.message - - def _wire_fields(self) -> dict[str, object]: - return {"error": self.error} - - -@dataclass(frozen=True, slots=True, kw_only=True) -class DuplexRawEvent(DuplexEvent): - """A session-internal event with no dedicated Realtime type (``duplex.``).""" - - wire_type = "duplex.raw" - - internal_type: str = "" - details: Mapping[str, object] = field(default_factory=dict) - - @property - def type(self) -> str: - return f"duplex.{self.internal_type}" - - def to_realtime(self) -> dict[str, object]: - return {"type": self.type, "event_id": self.event_id, "event": _wire_value(self.details)} - - -def error_event( - code: str, - message: str, - *, - event_id: object | None = None, - param: object | None = None, - extra: Mapping[str, object] | None = None, -) -> ErrorEvent: - """Build the error event for an internal code; ``event_id`` is the *client* event id.""" - return ErrorEvent( - code=code, - message=message, - related_event_id=event_id if isinstance(event_id, str) and event_id else None, - param=param if isinstance(param, str) and param else None, - extra=dict(extra or {}), - ) - - -#: Internal events that terminate a response/session and must never be dropped as stale. DOMAIN_TERMINAL_EVENTS = frozenset( { "response.done", @@ -682,7 +84,6 @@ DOMAIN_TERMINAL_EVENTS = frozenset( } ) -#: Internal streaming model-output events subject to epoch/close stale filtering. MODEL_OUTPUT_EVENTS = frozenset( { "response.created", @@ -705,13 +106,11 @@ MODEL_OUTPUT_EVENTS = frozenset( __all__ = [ - "DOMAIN_TERMINAL_EVENTS", - "MODEL_OUTPUT_EVENTS", - "REALTIME_ERROR_TYPES_BY_CODE", "AudioDelta", "AudioDone", "ContentPartAdded", "ContentPartDone", + "DOMAIN_TERMINAL_EVENTS", "DuplexEvent", "DuplexRawEvent", "ErrorEvent", @@ -727,11 +126,13 @@ __all__ = [ "ItemRetrieved", "ItemTruncated", "Listen", + "MODEL_OUTPUT_EVENTS", "OutputAudioCleared", "OutputItemAdded", "OutputItemDone", "OverlapDecision", "PlaybackAcknowledged", + "REALTIME_ERROR_TYPES_BY_CODE", "RateLimitsUpdated", "ResponseCreated", "ResponseDone", diff --git a/vllm_omni/engine/duplex/realtime_commands.py b/vllm_omni/engine/duplex/realtime_commands.py index 370364017..31b1c0949 100644 --- a/vllm_omni/engine/duplex/realtime_commands.py +++ b/vllm_omni/engine/duplex/realtime_commands.py @@ -3,26 +3,31 @@ """OpenAI Realtime client events -> :class:`DuplexCommand` (stateless). -This is the mapping half of the former ``entrypoints/duplex/realtime_input.py`` -translator. Everything that needed per-session state in the old translator -(input-buffer emptiness for commits, response-id fallbacks for cancels, -conversation-item lookups, VAD) is resolved by the session runner through the -helpers on :class:`~vllm_omni.engine.duplex.realtime_events.RealtimeProjectionState`; +This module is the *duplex binding* of the shared Realtime codec. Parsing, +validating and decoding a client event is not duplex-specific and lives in +``vllm_omni.protocol.realtime``, reached through ``vllm_omni.protocol.duplex``; +what lives here is the part that is: which +:class:`~vllm_omni.engine.duplex.commands.DuplexCommand` a decoded event +becomes, what the duplex engine can serve +(:data:`DUPLEX_REALTIME_CAPABILITIES`), and the mapping between a Realtime +output format and the duplex ``response_format`` vocabulary. + +Everything that needs per-session state (input-buffer emptiness for commits, +response-id fallbacks for cancels, conversation-item lookups, VAD) is resolved +by the session runner through the helpers on +:class:`~vllm_omni.engine.duplex.realtime_events.RealtimeProjectionState`; the commands produced here carry the raw client intent only. + +The shared names re-exported at the bottom are a compatibility surface for +existing importers. The canonical definitions are in +``vllm_omni.protocol.realtime`` and there is exactly one of each. """ from __future__ import annotations -import base64 -import binascii from collections.abc import Mapping -from dataclasses import dataclass, replace -from typing import Any, cast -from uuid import uuid4 +from typing import cast -import numpy as np - -from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate from vllm_omni.engine.duplex.commands import ( AckPlayback, AppendAudio, @@ -44,93 +49,58 @@ from vllm_omni.engine.duplex.commands import ( TruncateItem, UpdateSession, ) - -REALTIME_INPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "pcm_f32le", - "g711_ulaw", - "g711_alaw", -} -REALTIME_OUTPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "wav", - "pcm", - "g711_ulaw", - "g711_alaw", -} - -#: Wire hint keys copied verbatim from a client append onto the internal payload. -REALTIME_INPUT_HINT_KEYS = ( - "duration_ms", - "audio_duration_ms", - "audio_start_ms", - "audio_end_ms", - "is_speech", - "speech", - "speech_probability", - "vad", - "overlap_action", - "overlap", - "force_barge_in", - "force_listen", - "text", - "transcript", +from vllm_omni.protocol.duplex import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_INPUT_HINT_KEYS, + REALTIME_OUTPUT_AUDIO_FORMATS, + RealtimeInputDefaults, + RealtimeProtocolCapabilities, + RealtimeProtocolError, + apply_realtime_session_defaults, + copy_realtime_input_hints, + decode_audio_append, + input_audio_transcription_config, + input_explicitly_non_speech, + input_looks_like_speech, + input_transcript_from_item, + is_supported_realtime_input_format, + json_safe_realtime_payload, + normalize_conversation_item, + parse_realtime_audio_format, + realtime_audio_format_object, + realtime_max_output_tokens, + realtime_output_format, + realtime_overlap_fields, + text_chars_for_audio_ms_from_marks, + truncate_realtime_item_content, + validate_conversation_item_audio_formats, + validate_realtime_item_truncate, + validate_realtime_response_audio_formats, + validate_realtime_session_audio_formats, + validate_realtime_video_frames, + validate_session_payload, ) -@dataclass(frozen=True, slots=True) -class RealtimeInputDefaults: - """Session-level wire defaults an append may omit (derived from the session object).""" +def _validate_duplex_turn_detection(session_payload: Mapping[str, object]) -> str | None: + """The duplex engine's answer for ``turn_detection`` on a session object. - input_audio_format: str = "pcm16" - input_sample_rate_hz: int = 16000 - output_audio_format: str = "pcm16" - output_sample_rate_hz: int | None = None - overlap_silence_rms: float = 0.003 + Imported lazily: the validator lives with the Silero VAD backend, and a + connection that never sends a session object should not pay for loading it. + """ + from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection - def with_session_payload(self, session_payload: Mapping[str, object]) -> RealtimeInputDefaults: - """Return defaults updated from a Realtime ``session`` object (session.update).""" - values = apply_realtime_session_defaults(self, session_payload) - return values + return validate_realtime_turn_detection(session_payload) -# ---- audio format helpers ---- - - -def parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: - def normalize_format(fmt: str) -> str: - normalized = fmt.lower() - if normalized in {"audio/pcm", "pcm"}: - return "pcm16" - if normalized in {"audio/wav", "wav"}: - return "wav" - if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: - return "pcm16" - if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: - return "pcm_f32le" - if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: - return "g711_ulaw" - if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: - return "g711_alaw" - return fmt - - if isinstance(raw_format, str): - return normalize_format(raw_format), None - if not isinstance(raw_format, dict): - return raw_format, None - rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") - sample_rate_hz = int(rate) if isinstance(rate, int | float) and rate > 0 else None - fmt = raw_format.get("type") or raw_format.get("format") - if not isinstance(fmt, str): - return raw_format, sample_rate_hz - return normalize_format(fmt), sample_rate_hz +#: What the duplex engine can serve, for the shared session-object validator. +DUPLEX_REALTIME_CAPABILITIES = RealtimeProtocolCapabilities( + validate_turn_detection=_validate_duplex_turn_detection, +) def duplex_response_format(realtime_format: str) -> str: + """A Realtime output format -> the duplex ``response_format`` vocabulary.""" normalized = realtime_format.lower() if normalized in {"pcm16", "pcm_s16le", "s16le"}: return "pcm" @@ -141,433 +111,6 @@ def duplex_response_format(realtime_format: str) -> str: return "wav" -def realtime_output_format(duplex_format: object) -> str: - if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: - return duplex_format.lower() - if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": - return "pcm16" - return str(duplex_format or "wav") - - -def is_supported_realtime_input_format(fmt: object) -> bool: - return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS - - -def realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: - payload: dict[str, object] = {"type": "audio/pcm"} - elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": - payload = {"type": "audio/pcm_f32le"} - elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": - payload = {"type": "audio/g711_ulaw"} - elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": - payload = {"type": "audio/g711_alaw"} - else: - payload = {"type": "audio/wav"} - if sample_rate_hz is not None: - payload["rate"] = int(sample_rate_hz) - return payload - - -def validate_realtime_session_audio_formats(session_payload: Mapping[str, object]) -> str | None: - audio_config = session_payload.get("audio") - input_format: object = session_payload.get("input_audio_format") - if input_format is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - input_format = audio_input.get("format") - parsed_input, _ = parse_realtime_audio_format(input_format) - if input_format is not None and not ( - isinstance(parsed_input, str) and parsed_input.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - return f"Unsupported input_audio_format: {input_format}" - - output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - parsed_output, _ = parse_realtime_audio_format(output_format) - if output_format is not None and not ( - isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - return f"Unsupported output_audio_format: {output_format}" - return None - - -def validate_realtime_response_audio_formats(response_payload: Mapping[str, object]) -> str | None: - output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") - audio_config = response_payload.get("audio") - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - parsed_output, _ = parse_realtime_audio_format(output_format) - if output_format is not None and not ( - isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - return f"Unsupported output_audio_format: {output_format}" - return None - - -def validate_conversation_item_audio_formats(item: object) -> str | None: - if not isinstance(item, dict): - return None - content = item.get("content") - if not isinstance(content, list): - return None - for part in content: - if not isinstance(part, dict): - continue - if part.get("type") not in {"input_audio", "audio"}: - continue - raw_format = part.get("format") - parsed_format, _ = parse_realtime_audio_format(raw_format) - if raw_format is not None and not ( - isinstance(parsed_format, str) and parsed_format.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - return f"Unsupported input_audio format in conversation.item.create: {raw_format}" - return None - - -def apply_realtime_session_defaults( - defaults: RealtimeInputDefaults, - session_payload: Mapping[str, object], -) -> RealtimeInputDefaults: - """Derive the wire defaults a Realtime ``session`` object declares.""" - input_format: object = session_payload.get("input_audio_format") - audio_config = session_payload.get("audio") - if input_format is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - input_format = audio_input.get("format") - input_format, input_rate = parse_realtime_audio_format(input_format) - # Any: heterogeneous field values splatted into ``dataclasses.replace`` below. - updates: dict[str, Any] = {} - if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: - updates["input_audio_format"] = input_format - output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") - output_rate_raw: object | None = None - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") - output_format, output_rate = parse_realtime_audio_format(output_format) - if output_rate is None and isinstance(output_rate_raw, int | float) and output_rate_raw > 0: - output_rate = int(output_rate_raw) - if isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: - updates["output_audio_format"] = realtime_output_format(output_format) - sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") - if sample_rate is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") - if sample_rate is None: - sample_rate = input_rate - if isinstance(sample_rate, int | float) and sample_rate > 0: - updates["input_sample_rate_hz"] = int(sample_rate) - if isinstance(output_rate, int | float) and output_rate > 0: - updates["output_sample_rate_hz"] = int(output_rate) - overlap_fields = realtime_overlap_fields(session_payload) - overlap_silence_rms = overlap_fields.get("overlap_silence_rms") - if isinstance(overlap_silence_rms, int | float): - updates["overlap_silence_rms"] = max(0.0, float(overlap_silence_rms)) - return replace(defaults, **updates) if updates else defaults - - -def realtime_overlap_fields(session_payload: Mapping[str, object]) -> dict[str, object]: - fields: dict[str, object] = {} - if isinstance(session_payload.get("overlap_policy"), str): - fields["overlap_policy"] = session_payload["overlap_policy"] - for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): - value = session_payload.get(key) - if isinstance(value, int | float): - fields[key] = value - if isinstance(session_payload.get("playback_commit_policy"), str): - fields["playback_commit_policy"] = session_payload["playback_commit_policy"] - return fields - - -def json_safe_realtime_payload(payload: Mapping[str, object]) -> dict[str, object]: - clean: dict[str, object] = {} - for key, value in payload.items(): - if key == "extra_body": - continue - if isinstance(value, str | int | float | bool) or value is None: - clean[key] = value - elif isinstance(value, dict): - clean[key] = json_safe_realtime_payload(value) - elif isinstance(value, list): - clean[key] = [ - (json_safe_realtime_payload(item) if isinstance(item, dict) else item) - for item in value - if isinstance(item, str | int | float | bool | dict) or item is None - ] - return clean - - -def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: - transcription = session_payload.get("input_audio_transcription") - if isinstance(transcription, dict): - return transcription - audio_config = session_payload.get("audio") - if not isinstance(audio_config, dict): - return None - audio_input = audio_config.get("input") - if not isinstance(audio_input, dict): - return None - transcription = audio_input.get("transcription") - return transcription if isinstance(transcription, dict) else None - - -def realtime_max_output_tokens(value: object) -> int | None: - """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" - if value is None: - return None - if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: - return None - if isinstance(value, int) and value > 0: - return int(value) - return None - - -# ---- speech hints ---- - - -def input_explicitly_non_speech(event: Mapping[str, object]) -> bool: - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return not value - vad = event.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return not value - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) < 0.5 - probability = event.get("speech_probability") - return isinstance(probability, int | float) and float(probability) < 0.5 - - -def input_looks_like_speech( - event: Mapping[str, object], - *, - audio: object, - fmt: object, - overlap_silence_rms: float, -) -> bool: - if input_explicitly_non_speech(event): - return False - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return value - vad = event.get("vad") - if isinstance(vad, dict): - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) >= 0.5 - probability = event.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 - if fmt != "pcm_f32le": - return True - if isinstance(audio, bytes | bytearray): - raw = bytes(audio) - elif isinstance(audio, str): - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return True - else: - return True - if len(raw) < 4 or len(raw) % 4 != 0: - return True - samples = np.frombuffer(raw, dtype=np.float32) - if samples.size == 0: - return False - rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) - threshold = event.get("overlap_silence_rms") - if not isinstance(threshold, int | float): - vad = event.get("vad") - if isinstance(vad, dict): - threshold = vad.get("silence_rms") - silence_rms = float(threshold) if isinstance(threshold, int | float) else overlap_silence_rms - return rms >= max(0.0, silence_rms) - - -def copy_realtime_input_hints(source: Mapping[str, object], target: dict[str, object]) -> None: - for key in REALTIME_INPUT_HINT_KEYS: - if key in source: - target[key] = source[key] - - -def validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: - """Validate omni-duplex camera frames on input_audio_buffer.append. - - Wire contract matches the official MiniCPM-o duplex loop: one base - base64 JPEG per ~1 s audio chunk, optionally followed by that unit's - stacked composite tiling the sub-frames captured inside it (at most 2 - images either way). A caller-supplied ``max_slice_nums`` is rejected rather - than silently ignored: slicing is Stage 0's decision here, and Stage 0 - already applies the official HD suggestion for a stacked unit - (``max_slice_nums=[2, 1]``). The wire simply does not let the client choose - it. - """ - if max_slice_nums not in (None, 1): - return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" - if not isinstance(video_frames, list): - return "video_frames must be a list of base64-encoded images" - frames = [frame for frame in video_frames if frame is not None] - if len(frames) > 2: - return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" - for frame in frames: - if not isinstance(frame, str) or not frame: - return "video_frames entries must be non-empty base64 strings" - if len(frame) > 4_000_000: - return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" - try: - header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) - except (binascii.Error, ValueError): - return "video_frames entries must be valid base64" - if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): - return "video_frames entries must be JPEG or PNG images" - return None - - -# ---- conversation items ---- - - -def normalize_conversation_item(item: Mapping[str, object]) -> dict[str, object]: - normalized = dict(item) - normalized.setdefault("id", f"item_{uuid4().hex}") - normalized.setdefault("object", "realtime.item") - normalized.setdefault("type", "message") - normalized.setdefault("status", "completed") - if "role" not in normalized and normalized.get("type") == "message": - normalized["role"] = "user" - if not isinstance(normalized.get("content"), list): - normalized["content"] = [] - return normalized - - -def text_chars_for_audio_ms_from_marks( - audio_end_ms: int, - text_len: int, - marks: list[object], - *, - final_ms: object | None = None, -) -> int: - if text_len <= 0: - return 0 - clean_marks: list[tuple[int, int]] = [] - for mark in marks: - if not isinstance(mark, dict): - continue - raw_text_chars = mark.get("text_chars") - raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) - if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): - continue - clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) - if not clean_marks: - return 0 if audio_end_ms <= 0 else text_len - clean_marks.sort(key=lambda item: item[0]) - audio_end_ms = max(0, int(audio_end_ms)) - if audio_end_ms <= 0: - return 0 - previous_ms = 0 - previous_chars = 0 - for mark_ms, mark_chars in clean_marks: - mark_ms = max(previous_ms, mark_ms) - mark_chars = max(previous_chars, min(text_len, mark_chars)) - if audio_end_ms <= mark_ms: - if mark_ms <= previous_ms: - return mark_chars - ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) - return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) - previous_ms = mark_ms - previous_chars = mark_chars - if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: - if audio_end_ms >= int(final_ms): - return text_len - ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) - return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) - return text_len if audio_end_ms >= previous_ms else previous_chars - - -def truncate_realtime_item_content(item: dict[str, object], *, content_index: int, audio_end_ms: int) -> None: - content = item.get("content") - if not isinstance(content, list) or not content: - return - index = max(0, int(content_index)) - if index >= len(content): - return - part = content[index] - if not isinstance(part, dict): - return - transcript = part.get("transcript") - if not isinstance(transcript, str) or not transcript: - return - marks = part.get("audio_text_marks") - if isinstance(marks, list): - keep_chars = text_chars_for_audio_ms_from_marks( - audio_end_ms, - len(transcript), - marks, - final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), - ) - part["transcript"] = transcript[:keep_chars].rstrip() - return - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and duration_ms > 0: - keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) - part["transcript"] = transcript[:keep_chars].rstrip() - elif audio_end_ms <= 0: - part["transcript"] = "" - - -def validate_realtime_item_truncate(item: Mapping[str, object], *, content_index: int, audio_end_ms: int) -> str | None: - if item.get("type") != "message" or item.get("role") != "assistant": - return "conversation.item.truncate only supports assistant message items" - if audio_end_ms < 0: - return "conversation.item.truncate requires non-negative audio_end_ms" - content = item.get("content") - if not isinstance(content, list) or not content: - return None - index = max(0, int(content_index)) - if index >= len(content): - return f"conversation.item.truncate content_index out of range: {content_index}" - part = content[index] - if not isinstance(part, dict): - return "conversation.item.truncate target content part is invalid" - if part.get("type") not in {"audio", "output_audio"}: - return "conversation.item.truncate target content part must be audio" - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): - return "conversation.item.truncate audio_end_ms exceeds item audio duration" - return None - - -def input_transcript_from_item(item: Mapping[str, object]) -> str: - content = item.get("content") - if not isinstance(content, list): - return "" - parts: list[str] = [] - for part in content: - if not isinstance(part, dict): - continue - for key in ("transcript", "text"): - value = part.get(key) - if isinstance(value, str) and value: - if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: - parts.append(value) - break - return "".join(parts).strip() - - # ---- append / item audio conversion ---- @@ -577,67 +120,21 @@ def build_append_audio( defaults: RealtimeInputDefaults, hints_source: Mapping[str, object] | None = None, ) -> AppendAudio: - """Validate, convert (to 16 kHz ``pcm_f32le`` base64) and pack one audio append. - - Speech classification without VAD (``is_speech`` on the command) uses the - explicit wire hints and the RMS fallback exactly like the old translator. - """ - event_id = cast("str", event.get("event_id")) if isinstance(event.get("event_id"), str) else None - audio = event.get("audio") or event.get("delta") - fmt, format_rate = parse_realtime_audio_format( - event.get("format") or event.get("input_audio_format") or defaults.input_audio_format - ) - sample_rate_hz = ( - event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or defaults.input_sample_rate_hz - ) - if not is_supported_realtime_input_format(fmt): - raise DuplexCommandError( - f"Unsupported input_audio_format: {fmt}", - code="unsupported_audio_format", - event_id=event_id, - ) + """Decode one audio append (shared codec) and pack it as the duplex command.""" try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) - except ValueError as exc: - raise DuplexCommandError(str(exc), code="bad_event", event_id=event_id) from exc - hints: dict[str, object] = {} - if hints_source is not None: - copy_realtime_input_hints(hints_source, hints) - copy_realtime_input_hints(event, hints) - looks_like_speech = input_looks_like_speech( - {**(hints_source or {}), **event}, - audio=audio, - fmt=fmt, - overlap_silence_rms=defaults.overlap_silence_rms, - ) - video_frames: tuple[str, ...] = () - raw_frames = event.get("video_frames") - if raw_frames is not None: - frames_error = validate_realtime_video_frames(raw_frames, event.get("max_slice_nums")) - if frames_error is not None: - raise DuplexCommandError(frames_error, code="invalid_video_frames", event_id=event_id) - if isinstance(raw_frames, list): - video_frames = tuple(frame for frame in raw_frames if isinstance(frame, str) and frame) - duration_ms = hints.get("duration_ms", hints.get("audio_duration_ms")) - audio_end_ms = hints.get("audio_end_ms") - try: - audio_bytes = base64.b64decode(audio, validate=True) if isinstance(audio, str) and audio else b"" - except (binascii.Error, ValueError) as exc: - raise DuplexCommandError("input audio is not valid base64", code="bad_audio", event_id=event_id) from exc + decoded = decode_audio_append(event, defaults=defaults, hints_source=hints_source) + except RealtimeProtocolError as exc: + raise DuplexCommandError(str(exc), code=exc.code, event_id=exc.event_id) from exc return AppendAudio( - event_id=event_id, - audio=audio_bytes, - format=str(fmt), - sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, - is_speech=looks_like_speech, - video_frames=video_frames, - duration_ms=int(duration_ms) if isinstance(duration_ms, int | float) else None, - audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else None, - hints=hints, + event_id=decoded.event_id, + audio=decoded.audio, + format=decoded.format, + sample_rate_hz=decoded.sample_rate_hz, + is_speech=decoded.is_speech, + video_frames=decoded.video_frames, + duration_ms=decoded.duration_ms, + audio_end_ms=decoded.audio_end_ms, + hints=decoded.hints, ) @@ -664,14 +161,9 @@ def translate_realtime_command( if event_type == "session.update": session = payload.get("session") session_payload: Mapping[str, object] = session if isinstance(session, dict) else payload - format_error = validate_realtime_session_audio_formats(session_payload) - if format_error is not None: - raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) - from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection - - turn_detection_error = validate_realtime_turn_detection(session_payload) - if turn_detection_error is not None: - raise DuplexCommandError(turn_detection_error, code="unsupported_turn_detection", event_id=event_id) + rejection = validate_session_payload(session_payload, capabilities=DUPLEX_REALTIME_CAPABILITIES) + if rejection is not None: + raise DuplexCommandError(rejection.message, code=rejection.code, event_id=event_id) return UpdateSession(event_id=event_id, patch=dict(session_payload)) if event_type == "conversation.item.create": @@ -855,7 +347,10 @@ def translate_realtime_command( raise DuplexCommandError(f"Unknown duplex event type: {event_type}", code="unknown_event", event_id=event_id) +#: Names this module used to define, re-exported so existing importers keep +#: working. The canonical definitions live in ``vllm_omni.protocol.realtime``. __all__ = [ + "DUPLEX_REALTIME_CAPABILITIES", "REALTIME_INPUT_AUDIO_FORMATS", "REALTIME_INPUT_HINT_KEYS", "REALTIME_OUTPUT_AUDIO_FORMATS", diff --git a/vllm_omni/engine/duplex/realtime_events.py b/vllm_omni/engine/duplex/realtime_events.py index 91359f297..f67abbbc9 100644 --- a/vllm_omni/engine/duplex/realtime_events.py +++ b/vllm_omni/engine/duplex/realtime_events.py @@ -15,6 +15,12 @@ carries the input-side bookkeeping the old input translator kept (input buffer flags, conversation items, response-id fallbacks); the ``resolve_*`` / ``note_*`` helpers give the runner the same behaviour for the corresponding commands. + +What is *not* here is the model- and runtime-agnostic half of the codec --- +audio format negotiation, conversation-item shape and truncation, transcript +extraction, audio conversion. That lives in ``vllm_omni.protocol.realtime`` so +a non-duplex Realtime surface can use it without the duplex session; this +module is the duplex consumer of it (RFC #6592 P0a). """ from __future__ import annotations @@ -24,7 +30,6 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, cast from uuid import uuid4 -from vllm_omni.engine.duplex.audio import convert_output_audio from vllm_omni.engine.duplex.commands import ( AppendAudio, CancelResponse, @@ -74,10 +79,11 @@ from vllm_omni.engine.duplex.events import ( TranscriptDone, error_event, ) -from vllm_omni.engine.duplex.realtime_commands import ( +from vllm_omni.engine.duplex.realtime_commands import build_append_audio +from vllm_omni.protocol.duplex import ( RealtimeInputDefaults, apply_realtime_session_defaults, - build_append_audio, + convert_output_audio, copy_realtime_input_hints, input_looks_like_speech, input_transcript_from_item, diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index 7915500b7..0b8af49ab 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -28,7 +28,6 @@ from typing import TYPE_CHECKING, TypeVar from vllm.logger import init_logger -from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate from vllm_omni.engine.duplex.commands import ( AckPlayback, AppendAudio, @@ -103,6 +102,7 @@ from vllm_omni.engine.duplex.turn_detection import ( TurnDetectionResult, ) from vllm_omni.metrics.stats import StageRequestStats +from vllm_omni.protocol.duplex import convert_input_audio_with_rate if TYPE_CHECKING: from vllm.config import ModelConfig diff --git a/vllm_omni/entrypoints/duplex/realtime_input.py b/vllm_omni/entrypoints/duplex/realtime_input.py index be700ef62..79a081696 100644 --- a/vllm_omni/entrypoints/duplex/realtime_input.py +++ b/vllm_omni/entrypoints/duplex/realtime_input.py @@ -5,7 +5,9 @@ Everything that needed session state now lives engine-side (``vllm_omni.engine.duplex.realtime_commands`` for the command mapping, -``vllm_omni.engine.duplex.realtime_events`` for the projection state). What is +``vllm_omni.engine.duplex.realtime_events`` for the projection state), and the +model-agnostic parsing both of those build on lives in +``vllm_omni.protocol.realtime``. What is left here is the per-connection handshake policy (query-param defaults, autostart / resume-only rules, ``session.resume`` parsing), the wire defaults used to translate appends, and error rendering through the typed @@ -20,10 +22,8 @@ from typing import TYPE_CHECKING from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError from vllm_omni.engine.duplex.events import error_event -from vllm_omni.engine.duplex.realtime_commands import ( - RealtimeInputDefaults, - translate_realtime_command, -) +from vllm_omni.engine.duplex.realtime_commands import translate_realtime_command +from vllm_omni.protocol.duplex import RealtimeInputDefaults if TYPE_CHECKING: from starlette.websockets import WebSocket diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index 364f715c0..ff1abbe25 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -36,7 +36,6 @@ from vllm_omni.engine.duplex.events import ( SessionResyncRequired, ) from vllm_omni.engine.duplex.messages import DuplexSessionError -from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request from vllm_omni.entrypoints.duplex.session_attachment import ( DuplexJournalGapError, @@ -52,6 +51,7 @@ from vllm_omni.entrypoints.duplex.websocket import ( receive_text_with_timeout, ) from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle +from vllm_omni.protocol.duplex import RealtimeInputDefaults logger = init_logger(__name__) diff --git a/vllm_omni/protocol/__init__.py b/vllm_omni/protocol/__init__.py new file mode 100644 index 000000000..e241a7924 --- /dev/null +++ b/vllm_omni/protocol/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Wire protocols vLLM-Omni speaks, independent of who serves or runs them. + +A module under ``vllm_omni.protocol`` describes *what is on the wire*: how a +client payload is parsed and validated, and what shape a server payload has. +It never decides how a request is executed, which model answers it, or where +the session state lives. That keeps one protocol implementation usable by more +than one runtime (see ``vllm_omni.protocol.realtime``). + +Dependency rule, asserted by ``tests/protocol/realtime/test_protocol_import_boundary.py``: +nothing here may import ``vllm_omni.engine``, ``vllm_omni.entrypoints``, +``vllm_omni.model_executor``, ``vllm_omni.worker`` or ``vllm_omni.clients``. +The arrows point inwards --- those packages import the protocol, never the +other way round. +""" diff --git a/vllm_omni/protocol/duplex/__init__.py b/vllm_omni/protocol/duplex/__init__.py new file mode 100644 index 000000000..d89a01322 --- /dev/null +++ b/vllm_omni/protocol/duplex/__init__.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""vLLM-Omni's full-duplex extension of the OpenAI Realtime wire protocol. + +OpenAI's Realtime API assumes one turn at a time: the client commits input, the +server answers. A full-duplex session does not --- the model listens while it +speaks, it decides when to take a turn, the client reports playback progress, +and a dropped socket re-attaches to a session that kept running. Those need +events and commands the GA vocabulary does not have. + +This package is that delta. It builds on ``vllm_omni.protocol.realtime`` (same +base classes, same pure wire rendering, no duplication) and, like it, owns no +session and imports no runtime: the duplex engine depends on this package, not +the other way round. + +The one door +------------ +It is also the *only* door a duplex consumer uses. ``protocol.duplex.commands`` +and ``protocol.duplex.events`` carry the complete vocabulary --- Tier 1 classes +re-exported unchanged alongside the Tier 2 and Tier 3 ones --- and this module +re-exports the Tier 1 helper functions. So the layering is a chain, not a mesh:: + + protocol.realtime (OpenAI's surface) + ^ + | imports + | + protocol.duplex (this package: re-exports Tier 1, adds Tier 2 and 3) + ^ + | imports + | + duplex engine / entrypoints / clients + +``tests/protocol/duplex/test_duplex_protocol_facade.py`` asserts that the last +arrow is the only one: no duplex consumer imports ``protocol.realtime`` +directly. The payoff is that a helper which later needs a duplex-specific +version --- ``convert_input_audio_with_rate`` is the standing example, it +resamples to MiniCPM-o's 16 kHz rather than the client's rate --- is overridden +in one file instead of at every call site. + +What stays outside +------------------ +The *mailbox* half of a command does not live here. A duplex command carries +two different things: the client event it decodes from (wire, this package) and +the dictionary the session runner consumes (engine). They genuinely differ --- +``session.update``, ``conversation.item.create`` / ``.delete`` / ``.truncate`` +all travel on the runner's ``turn.signal`` channel --- so ``payload()`` and its +``type`` stay in ``vllm_omni.engine.duplex.commands``. Events have no such +half, which is why :data:`~vllm_omni.protocol.duplex.events.DuplexEvent` can be +a plain alias of ``RealtimeEvent`` while ``DuplexCommand`` cannot. +""" + +from vllm_omni.protocol.duplex.errors import ( + REALTIME_ERROR_TYPES_BY_CODE, + RealtimeProtocolError, + realtime_error_type, +) +from vllm_omni.protocol.realtime.audio import ( + MAX_INPUT_SAMPLE_RATE_HZ, + MIN_INPUT_SAMPLE_RATE_HZ, + convert_input_audio_with_rate, + convert_output_audio, + decode_g711_alaw, + decode_g711_ulaw, + encode_float32_mono_wav_base64, + encode_g711_alaw, + encode_g711_ulaw, + resample_pcm16_mono, + validate_input_sample_rate_hz, + wav_payload_to_pcm16, +) +from vllm_omni.protocol.realtime.audio_input import ( + REALTIME_INPUT_HINT_KEYS, + RealtimeAudioAppend, + copy_realtime_input_hints, + decode_audio_append, + input_explicitly_non_speech, + input_looks_like_speech, +) +from vllm_omni.protocol.realtime.capabilities import ( + RealtimeProtocolCapabilities, + RealtimeSessionRejection, + validate_session_payload, +) +from vllm_omni.protocol.realtime.formats import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_OUTPUT_AUDIO_FORMATS, + is_supported_realtime_input_format, + parse_realtime_audio_format, + realtime_audio_format_object, + realtime_output_format, + validate_conversation_item_audio_formats, + validate_realtime_response_audio_formats, + validate_realtime_session_audio_formats, +) +from vllm_omni.protocol.realtime.items import ( + input_transcript_from_item, + normalize_conversation_item, + text_chars_for_audio_ms_from_marks, + truncate_realtime_item_content, + validate_realtime_item_truncate, + validate_realtime_video_frames, +) +from vllm_omni.protocol.realtime.session import ( + RealtimeInputDefaults, + apply_realtime_session_defaults, + input_audio_transcription_config, + json_safe_realtime_payload, + realtime_max_output_tokens, + realtime_overlap_fields, +) + +__all__ = [ + "MAX_INPUT_SAMPLE_RATE_HZ", + "MIN_INPUT_SAMPLE_RATE_HZ", + "REALTIME_ERROR_TYPES_BY_CODE", + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_INPUT_HINT_KEYS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "RealtimeAudioAppend", + "RealtimeInputDefaults", + "RealtimeProtocolCapabilities", + "RealtimeProtocolError", + "RealtimeSessionRejection", + "apply_realtime_session_defaults", + "convert_input_audio_with_rate", + "convert_output_audio", + "copy_realtime_input_hints", + "decode_audio_append", + "decode_g711_alaw", + "decode_g711_ulaw", + "encode_float32_mono_wav_base64", + "encode_g711_alaw", + "encode_g711_ulaw", + "input_audio_transcription_config", + "input_explicitly_non_speech", + "input_looks_like_speech", + "input_transcript_from_item", + "is_supported_realtime_input_format", + "json_safe_realtime_payload", + "normalize_conversation_item", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_error_type", + "realtime_max_output_tokens", + "realtime_output_format", + "realtime_overlap_fields", + "resample_pcm16_mono", + "text_chars_for_audio_ms_from_marks", + "truncate_realtime_item_content", + "validate_conversation_item_audio_formats", + "validate_input_sample_rate_hz", + "validate_realtime_item_truncate", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", + "validate_realtime_video_frames", + "validate_session_payload", + "wav_payload_to_pcm16", +] diff --git a/vllm_omni/protocol/duplex/commands.py b/vllm_omni/protocol/duplex/commands.py new file mode 100644 index 000000000..1dbb9bf25 --- /dev/null +++ b/vllm_omni/protocol/duplex/commands.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The client-event vocabulary a full-duplex session accepts. + +Seven of the seventeen are events OpenAI's Realtime API has no equivalent for, +because they only make sense when the model and the user can talk at the same +time: the user cuts in (``barge_in``), the client reports how much audio it +played (``playback.ack``), a turn transition is signalled locally +(``turn.signal``), input is abandoned without committing (``input.cancel``), +text is pushed into a live turn (``input.text.append``), and the session is +kept alive or closed explicitly (``session.heartbeat`` / ``session.close``). +Two more are OpenAI commands carrying duplex-only fields (Tier 2). + +The remaining eight are OpenAI's, unchanged, and re-exported from +``vllm_omni.protocol.realtime.commands`` rather than redeclared --- they are +the same objects. They are re-exported so this module is the *whole* +vocabulary and a duplex consumer never has to import the Tier 1 package +directly; keeping them declared there means a GA-only consumer is still not +handed events its clients never send. + +None of these carries the engine's mailbox rendering: ``payload()`` and its +channel live in ``vllm_omni.engine.duplex.commands``, because for four of the +seventeen the runner's channel is not the client event type at all. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import ClassVar + +from vllm_omni.protocol.realtime import commands as realtime_commands +from vllm_omni.protocol.realtime.commands import ( + CancelResponse, + ClearInput, + ClearOutputAudio, + CreateItem, + CreateResponse, + DeleteItem, + RealtimeCommand, + TruncateItem, + UpdateSession, +) + +# ---- Tier 2: OpenAI names carrying vLLM-Omni extensions ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendAudio(realtime_commands.AppendAudio): + """Tier 2: OpenAI's ``input_audio_buffer.append`` plus the duplex hints. + + OpenAI's append is just base64 audio. The duplex lane also lets a client + declare this chunk's own format and rate, whether it believes the chunk is + speech, camera frames captured alongside it, and its timing --- all + additive, all ignored by a stock client. + """ + + #: Model-neutral hints carried through from the wire (rms, vad, transcript hints ...). + format: str = "pcm16" + sample_rate_hz: int | None = None + is_speech: bool | None = None + video_frames: tuple[str, ...] = () + duration_ms: int | None = None + audio_end_ms: int | None = None + hints: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Commit(realtime_commands.Commit): + """Tier 2: OpenAI's ``input_audio_buffer.commit`` plus duplex turn control. + + OpenAI's commit carries nothing and always produces an item. Here a commit + may be non-final, may decline to start a response, and may declare itself + silence --- the semantic divergence documented in + ``docs/serving/realtime_duplex_api.md`` ("Commit != response"). + """ + + #: Realtime conversation item created for this commit (wire correlation only). + #: ``None`` means "no explicit request": the runner decides on commit (auto-response + #: sessions answer on their own); ``True`` / ``False`` force it. + create_response: bool | None = None + is_speech: bool | None = None + realtime_item_id: str | None = None + + +# ---- Tier 3: vLLM-Omni only ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendText(RealtimeCommand): + wire_type: ClassVar[str] = "input.text.append" + text: str + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelInput(RealtimeCommand): + wire_type: ClassVar[str] = "input.cancel" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BargeIn(RealtimeCommand): + wire_type: ClassVar[str] = "barge_in" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SignalTurn(RealtimeCommand): + """Generic ``turn.signal`` (local turn transitions such as ``user_started``).""" + + wire_type: ClassVar[str] = "turn.signal" + event: str + signal_payload: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AckPlayback(RealtimeCommand): + wire_type: ClassVar[str] = "playback.ack" + played_ms: int + committed_ms: int | None = None + response_id: str | None = None + item_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Heartbeat(RealtimeCommand): + wire_type: ClassVar[str] = "session.heartbeat" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CloseSession(RealtimeCommand): + """Graceful close requested through the command stream (``session.close``).""" + + wire_type: ClassVar[str] = "session.close" + reason: str = "client_close" + + +#: The complete command vocabulary a duplex client may send: the eight Tier 1 +#: classes re-exported unchanged, the two Tier 2 classes defined above, and the +#: seven Tier 3 ones. A duplex consumer imports from here and never reaches +#: past this module into ``vllm_omni.protocol.realtime``, so a command that +#: later grows a duplex extension changes only this file. +__all__ = [ + # Tier 1 --- re-exported from ``vllm_omni.protocol.realtime.commands``. + "CancelResponse", + "ClearInput", + "ClearOutputAudio", + "CreateItem", + "CreateResponse", + "DeleteItem", + "RealtimeCommand", + "TruncateItem", + "UpdateSession", + # Tier 2 --- an OpenAI command plus duplex-only fields. + "AppendAudio", + "Commit", + # Tier 3 --- vLLM-Omni only. + "AckPlayback", + "AppendText", + "BargeIn", + "CancelInput", + "CloseSession", + "Heartbeat", + "SignalTurn", +] diff --git a/vllm_omni/protocol/duplex/errors.py b/vllm_omni/protocol/duplex/errors.py new file mode 100644 index 000000000..418c98712 --- /dev/null +++ b/vllm_omni/protocol/duplex/errors.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""vLLM-Omni's Realtime error-code vocabulary. + +OpenAI's Realtime API standardises the ``error.type`` classes +(``invalid_request_error`` / ``server_error`` / ``rate_limit_error``) but not +the ``code`` values inside them. Ours are listed here --- Tier 3 in +``docs/serving/realtime_duplex_api.md``, "the full error-code vocabulary" --- +so ``protocol/realtime`` stays the OpenAI surface and nothing more. +""" + +from __future__ import annotations + +from vllm_omni.protocol.realtime.errors import RealtimeProtocolError + +__all__ = ["REALTIME_ERROR_TYPES_BY_CODE", "RealtimeProtocolError", "realtime_error_type"] + + +#: OpenAI Realtime ``error.type`` for each internal error code. +REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = { + "bad_event": "invalid_request_error", + "bad_audio": "invalid_request_error", + "config_timeout": "invalid_request_error", + "invalid_json": "invalid_request_error", + "event_too_large": "invalid_request_error", + "unknown_event": "invalid_request_error", + "internal_error": "server_error", + "runtime_append_failed": "server_error", + "runtime_append_task_failed": "server_error", + "runtime_signal_failed": "server_error", + "runtime_abort_failed": "server_error", + "runtime_data_plane_stream_failed": "server_error", + "runtime_data_plane_text_without_audio": "server_error", + "resource_exhausted": "rate_limit_error", + "session_exists": "invalid_request_error", + "session_closed": "invalid_request_error", + "unknown_session": "invalid_request_error", + "invalid_duplex_runtime_config": "invalid_request_error", + "instructions_update_unsupported": "invalid_request_error", + "persona_update_unsupported": "invalid_request_error", + "voice_update_unsupported": "invalid_request_error", + "unsupported_nemotron_duplex_mode": "invalid_request_error", + "unsupported_native_response_options": "invalid_request_error", + "runtime_touch_failed": "server_error", + "engine_error": "server_error", + "input_backpressure": "rate_limit_error", + "response_already_active": "invalid_request_error", + "response_not_active": "invalid_request_error", + "response_create_without_input": "invalid_request_error", + "text_only_turn_unsupported": "invalid_request_error", + "commit_aborted": "server_error", + "input_audio_buffer_empty": "invalid_request_error", + "missing_item_id": "invalid_request_error", + "item_not_found": "invalid_request_error", + "playback_item_mismatch": "invalid_request_error", + "playback_item_not_found": "invalid_request_error", + "playback_ack_too_late": "invalid_request_error", + "unsupported_audio_format": "invalid_request_error", + "unsupported_turn_detection": "invalid_request_error", + "unsupported_ref_audio_path": "invalid_request_error", + "ref_audio_required": "invalid_request_error", + "model_update_unsupported": "invalid_request_error", + "voice_update_after_audio_unsupported": "invalid_request_error", + "ref_audio_update_unsupported": "invalid_request_error", + "native_text_append_unsupported": "invalid_request_error", + "invalid_video_frames": "invalid_request_error", + "invalid_function_call_output": "invalid_request_error", + "server_vad_unavailable": "server_error", +} + + +def realtime_error_type(code: str) -> str: + """The OpenAI ``error.type`` bucket for an internal error code.""" + return REALTIME_ERROR_TYPES_BY_CODE.get(code, "invalid_request_error") diff --git a/vllm_omni/protocol/duplex/events.py b/vllm_omni/protocol/duplex/events.py new file mode 100644 index 000000000..f2abe748e --- /dev/null +++ b/vllm_omni/protocol/duplex/events.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""vLLM-Omni's half of the Realtime server-event vocabulary. + +``docs/serving/realtime_duplex_api.md`` sorts every message into three tiers. +Tier 1 --- identical to OpenAI --- is ``vllm_omni.protocol.realtime.events``. +The other two are here: + +**Tier 2, OpenAI names carrying vLLM-Omni extensions.** Same wire type, same +OpenAI fields, plus additive keys a stock client ignores. Each one subclasses +its Tier 1 twin and declares only what it adds, so the OpenAI surface stays +honest: a consumer that wants pure GA takes ``protocol/realtime`` and gets +exactly OpenAI's fields. + +**Tier 3, vLLM-Omni only.** Events OpenAI has no equivalent for, because they +only make sense when the model and the user can talk at the same time: the +model decides when to listen and when to speak, the client reports how much +audio it played, overlapping speech needs a recorded decision, and a dropped +socket can re-attach to a live session. + +Same-named classes are deliberate --- ``duplex.events.SessionCreated`` *is* +``realtime.events.SessionCreated`` with our extensions. Import whichever +surface you mean. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import cast + +from vllm_omni.protocol.duplex.errors import REALTIME_ERROR_TYPES_BY_CODE +from vllm_omni.protocol.realtime import events as realtime_events +from vllm_omni.protocol.realtime.events import ( + AudioDone, + ContentPartAdded, + ContentPartDone, + FunctionCallArgumentsDelta, + FunctionCallArgumentsDone, + InputCleared, + InputTranscriptionCompleted, + ItemAdded, + ItemCreated, + ItemDone, + ItemRetrieved, + OutputAudioCleared, + OutputItemAdded, + OutputItemDone, + RateLimitsUpdated, + RealtimeEvent, + ResponseScopedEvent, + SessionUpdated, + SpeechStarted, + SpeechStopped, + TextDelta, + TextDone, + TranscriptDelta, + TranscriptDone, + new_event_id, + wire_value, +) + +#: The duplex event base class *is* the Realtime one: unlike a command, an +#: event has no engine-internal half to carry --- what the session emits is +#: exactly what the socket sends. Named here rather than in the engine so the +#: duplex vocabulary is described in one place, and kept as an alias rather +#: than an empty subclass so ``isinstance`` and the class identity tests in +#: ``tests/protocol/`` stay meaningful. +DuplexEvent = RealtimeEvent + +# ---- Tier 2: OpenAI names carrying vLLM-Omni extensions ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionCreated(realtime_events.SessionCreated): + """Tier 2: OpenAI's ``SessionCreated`` plus the resume credentials a reconnecting client needs.""" + + #: Transport resume credentials (set by the websocket handler when resume is supported). + optional_wire_fields = frozenset({"attachment_generation", "resume_token"}) + attachment_generation: int | None = None + resume_token: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseCreated(realtime_events.ResponseCreated): + """Tier 2: OpenAI's ``ResponseCreated`` plus a top-level ``response_id``. + + OpenAI carries the id nested as ``response.id``; the duplex lane also + reports it at the top level, where every other response-scoped event + already has one. + """ + + response_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseDone(realtime_events.ResponseDone): + """Tier 2: OpenAI's ``ResponseDone`` plus a top-level ``response_id``. + + OpenAI carries the id nested as ``response.id``; the duplex lane also + reports it at the top level, where every other response-scoped event + already has one. + """ + + response_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDelta(realtime_events.AudioDelta): + """Tier 2: OpenAI's ``AudioDelta`` plus the chunk's own format/rate and the duplex playback metadata.""" + + optional_wire_fields = frozenset({"sample_rate_hz", "metadata"}) + format: str = "pcm16" + sample_rate_hz: int | None = None + metadata: Mapping[str, object] | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCommitted(realtime_events.InputCommitted): + """Tier 2: OpenAI's ``InputCommitted`` plus the native ``input.committed`` event echoed under ``event``.""" + + #: The session-internal ``input.committed`` event. + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + # Explicit parent call, not ``super()``: ``@dataclass(slots=True)`` + # rebuilds the class, so the zero-argument form's ``__class__`` cell + # points at the pre-slots class and is no longer in the MRO. + base = realtime_events.__dict__[type(self).__name__]._wire_fields(self) + return {**base, "event": wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDeleted(realtime_events.ItemDeleted): + """Tier 2: OpenAI's ``ItemDeleted`` plus the originating event echoed under ``event``.""" + + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + # Explicit parent call, not ``super()``: ``@dataclass(slots=True)`` + # rebuilds the class, so the zero-argument form's ``__class__`` cell + # points at the pre-slots class and is no longer in the MRO. + base = realtime_events.__dict__[type(self).__name__]._wire_fields(self) + return {**base, "event": wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemTruncated(realtime_events.ItemTruncated): + """Tier 2: OpenAI's ``ItemTruncated`` plus the originating event echoed under ``event``.""" + + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + # Explicit parent call, not ``super()``: ``@dataclass(slots=True)`` + # rebuilds the class, so the zero-argument form's ``__class__`` cell + # points at the pre-slots class and is no longer in the MRO. + base = realtime_events.__dict__[type(self).__name__]._wire_fields(self) + return {**base, "event": wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ErrorEvent(realtime_events.ErrorEvent): + """Tier 2: OpenAI's ``ErrorEvent`` plus vLLM-Omni's error codes and any extra keys merged into the envelope.""" + + #: Extra keys merged into the wire ``error`` object (``retryable`` ...). + extra: Mapping[str, object] = field(default_factory=dict) + + @property + def error_type(self) -> str: + return REALTIME_ERROR_TYPES_BY_CODE.get(self.code, "invalid_request_error") + + @property + def error(self) -> dict[str, object]: + # Explicit parent call, not ``super()``: ``@dataclass(slots=True)`` + # rebuilds the class, so the zero-argument form's ``__class__`` cell + # points at the pre-slots class and is no longer in the MRO. Reached + # through ``__dict__`` because attribute access on the class gives the + # getter function, not the ``property`` that owns it. + error = dict(realtime_events.ErrorEvent.__dict__["error"].fget(self)) + error.update(cast("Mapping[str, object]", wire_value(self.extra))) + return error + + +# ---- Tier 3: vLLM-Omni only ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResumed(SessionCreated): + wire_type = "session.resumed" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **RealtimeEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionReplaced(RealtimeEvent): + """A newer connection took the session over (sent to the replaced socket).""" + + wire_type = "session.replaced" + + attachment_generation: int = 0 + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **RealtimeEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResyncRequired(RealtimeEvent): + """The replay journal cannot bridge the gap; the client must start over.""" + + wire_type = "session.resync_required" + + reason: str = "journal_gap" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **RealtimeEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionClosed(RealtimeEvent): + wire_type = "session.closed" + + reason: str = "closed" + #: The session-internal close event (kept for clients that inspect it). + details: Mapping[str, object] = field(default_factory=dict) + + @property + def is_terminal(self) -> bool: + return True + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason, "event": wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionExpired(SessionClosed): + wire_type = "session.expired" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionHeartbeatAck(RealtimeEvent): + wire_type = "session.heartbeat_ack" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TurnEvent(RealtimeEvent): + """Turn-state transition (``turn.event``).""" + + wire_type = "turn.event" + + event: str = "" + turn_state: str = "" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "event": self.event, "turn_state": self.turn_state, "epoch": self.epoch} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Listen(RealtimeEvent): + """The model decided to keep listening (no spoken response for this turn).""" + + wire_type = "response.listen" + optional_wire_fields = frozenset({"response_id"}) + + response_id: str | None = None # type: ignore[assignment] + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + response: dict[str, object] = { + "object": "realtime.response", + "status": "listening", + "metadata": wire_value(self.details), + } + data: dict[str, object] = {"session_id": self.session_id, "epoch": self.epoch, "response": response} + if self.response_id: + response["id"] = self.response_id + data["response_id"] = self.response_id + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Speak(ResponseScopedEvent): + wire_type = "response.speak" + + metadata: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OverlapDecision(RealtimeEvent): + wire_type = "overlap.decision" + + policy: str | None = None + action: str | None = None + reason: str | None = None + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "session_id": self.session_id, + "epoch": self.epoch, + "policy": self.policy, + "action": self.action, + "reason": self.reason, + "metadata": wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PlaybackAcknowledged(RealtimeEvent): + wire_type = "playback.acknowledged" + + #: The session-internal ``playback.acknowledged`` event (cursor, committed_ms ...). + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return {"event": wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexRawEvent(RealtimeEvent): + """A session-internal event with no dedicated Realtime type (``duplex.``).""" + + wire_type = "duplex.raw" + + internal_type: str = "" + details: Mapping[str, object] = field(default_factory=dict) + + @property + def type(self) -> str: + return f"duplex.{self.internal_type}" + + def to_realtime(self) -> dict[str, object]: + return {"type": self.type, "event_id": self.event_id, "event": wire_value(self.details)} + + +def error_event( + code: str, + message: str, + *, + event_id: object | None = None, + param: object | None = None, + extra: Mapping[str, object] | None = None, +) -> ErrorEvent: + """Build the error event for an internal code; ``event_id`` is the *client* event id.""" + return ErrorEvent( + code=code, + message=message, + related_event_id=event_id if isinstance(event_id, str) and event_id else None, + param=param if isinstance(param, str) and param else None, + extra=dict(extra or {}), + ) + + +#: The complete event vocabulary a duplex client may receive: 22 Tier 1 +#: classes re-exported unchanged, the 8 Tier 2 ones defined above, and 12 +#: Tier 3 ones. A duplex consumer imports from here and never reaches past +#: this module into ``vllm_omni.protocol.realtime``. +__all__ = [ + # Tier 1 --- re-exported from ``vllm_omni.protocol.realtime.events``. + "AudioDone", + "ContentPartAdded", + "ContentPartDone", + "FunctionCallArgumentsDelta", + "FunctionCallArgumentsDone", + "InputCleared", + "InputTranscriptionCompleted", + "ItemAdded", + "ItemCreated", + "ItemDone", + "ItemRetrieved", + "OutputAudioCleared", + "OutputItemAdded", + "OutputItemDone", + "RateLimitsUpdated", + "SessionUpdated", + "SpeechStarted", + "SpeechStopped", + "TextDelta", + "TextDone", + "TranscriptDelta", + "TranscriptDone", + # Tier 2 --- an OpenAI event plus duplex-only fields. + "AudioDelta", + "ErrorEvent", + "InputCommitted", + "ItemDeleted", + "ItemTruncated", + "ResponseCreated", + "ResponseDone", + "SessionCreated", + # Tier 3 --- vLLM-Omni only. + "DuplexRawEvent", + "Listen", + "OverlapDecision", + "PlaybackAcknowledged", + "SessionClosed", + "SessionExpired", + "SessionHeartbeatAck", + "SessionReplaced", + "SessionResumed", + "SessionResyncRequired", + "Speak", + "TurnEvent", + # Bases and helpers shared with the Tier 1 surface. + "DuplexEvent", + "REALTIME_ERROR_TYPES_BY_CODE", + "RealtimeEvent", + "ResponseScopedEvent", + "error_event", + "new_event_id", + "wire_value", +] diff --git a/vllm_omni/protocol/realtime/__init__.py b/vllm_omni/protocol/realtime/__init__.py new file mode 100644 index 000000000..af7194cb5 --- /dev/null +++ b/vllm_omni/protocol/realtime/__init__.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The OpenAI Realtime wire codec, shared by every vLLM-Omni Realtime surface. + +What this is +------------ +The model-agnostic half of ``WS /v1/realtime``: how a client event is parsed, +which audio formats and session fields are accepted, what a conversation item +means, how audio is decoded and resampled, and the shape of the error envelope +(the *codes* inside it are a consumer's own --- ours are in +``vllm_omni.protocol.duplex.errors``). It is a library of pure functions and value objects. + +What it is not +-------------- +It owns no session. It does not decide when a response starts, what a model is +prompted with, whether a turn is over, or where committed audio goes. Those are +runtime decisions, and after PR #7413 a duplex session's runtime decisions all +belong to ``vllm_omni.engine.duplex`` --- the session runner, the typed +``DuplexCommand`` / ``DuplexEvent`` boundary and the ``DuplexModelPlugin`` seam. +This package sits *under* that: the duplex engine binds the codec to its own +command and event vocabulary in +``vllm_omni.engine.duplex.realtime_commands`` and +``vllm_omni.engine.duplex.realtime_events``. + +Why it is separate (RFC #6592, P0a) +----------------------------------- +So that a second Realtime surface --- a Qwen3-Omni GA profile, say --- can reuse +the parsing, validation, format negotiation and error envelope without adopting +the duplex session control plane (lease manager, commit policy, model channel), +and without a second copy of the codec drifting away from this one. A consumer +states what it can serve through +:class:`~vllm_omni.protocol.realtime.capabilities.RealtimeProtocolCapabilities` +and keeps its own session state. + +Known duplex-shaped defaults +---------------------------- +This package is model-agnostic in its *dependencies*, not yet in every +*default*. Two constants still carry the assumptions of its first consumer and +a second one will have to parameterize them: + +* :func:`~vllm_omni.protocol.realtime.audio.convert_input_audio_with_rate` + resamples to 16 kHz, which is MiniCPM-o's feature-extractor rate. A GA client + sending 24 kHz is converted down rather than served at its own rate. +* :func:`~vllm_omni.protocol.realtime.items.validate_realtime_video_frames` + encodes the MiniCPM-o duplex camera contract (at most two frames per append, + ``max_slice_nums`` not selectable on the wire). + +Both are behaviour this extraction deliberately did **not** change (RFC #6592 +P0a is behaviour-preserving); making them per-consumer is P0b/P1 work. + +Dependency rule +--------------- +Nothing here may import ``vllm_omni.engine``, ``vllm_omni.entrypoints``, +``vllm_omni.model_executor``, ``vllm_omni.worker`` or ``vllm_omni.clients``; +``tests/protocol/realtime/test_protocol_import_boundary.py`` asserts it. +""" + +from vllm_omni.protocol.realtime.audio import ( + MAX_INPUT_SAMPLE_RATE_HZ, + MIN_INPUT_SAMPLE_RATE_HZ, + convert_input_audio_with_rate, + convert_output_audio, + encode_float32_mono_wav_base64, + resample_pcm16_mono, + validate_input_sample_rate_hz, + wav_payload_to_pcm16, +) +from vllm_omni.protocol.realtime.audio_input import ( + REALTIME_INPUT_HINT_KEYS, + RealtimeAudioAppend, + copy_realtime_input_hints, + decode_audio_append, + input_explicitly_non_speech, + input_looks_like_speech, +) +from vllm_omni.protocol.realtime.capabilities import ( + RealtimeProtocolCapabilities, + RealtimeSessionRejection, + validate_session_payload, +) +from vllm_omni.protocol.realtime.commands import RealtimeCommand +from vllm_omni.protocol.realtime.errors import RealtimeProtocolError +from vllm_omni.protocol.realtime.events import ( + RealtimeEvent, + ResponseScopedEvent, + new_event_id, + wire_value, +) +from vllm_omni.protocol.realtime.formats import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_OUTPUT_AUDIO_FORMATS, + is_supported_realtime_input_format, + parse_realtime_audio_format, + realtime_audio_format_object, + realtime_output_format, + validate_conversation_item_audio_formats, + validate_realtime_response_audio_formats, + validate_realtime_session_audio_formats, +) +from vllm_omni.protocol.realtime.items import ( + input_transcript_from_item, + normalize_conversation_item, + text_chars_for_audio_ms_from_marks, + truncate_realtime_item_content, + validate_realtime_item_truncate, + validate_realtime_video_frames, +) +from vllm_omni.protocol.realtime.session import ( + RealtimeInputDefaults, + apply_realtime_session_defaults, + input_audio_transcription_config, + json_safe_realtime_payload, + realtime_max_output_tokens, + realtime_overlap_fields, +) + +__all__ = [ + "MAX_INPUT_SAMPLE_RATE_HZ", + "MIN_INPUT_SAMPLE_RATE_HZ", + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_INPUT_HINT_KEYS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "RealtimeAudioAppend", + "RealtimeCommand", + "RealtimeEvent", + "RealtimeInputDefaults", + "RealtimeProtocolCapabilities", + "RealtimeProtocolError", + "RealtimeSessionRejection", + "ResponseScopedEvent", + "apply_realtime_session_defaults", + "convert_input_audio_with_rate", + "convert_output_audio", + "copy_realtime_input_hints", + "decode_audio_append", + "encode_float32_mono_wav_base64", + "input_audio_transcription_config", + "input_explicitly_non_speech", + "input_looks_like_speech", + "input_transcript_from_item", + "is_supported_realtime_input_format", + "json_safe_realtime_payload", + "new_event_id", + "normalize_conversation_item", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_max_output_tokens", + "realtime_output_format", + "realtime_overlap_fields", + "resample_pcm16_mono", + "text_chars_for_audio_ms_from_marks", + "truncate_realtime_item_content", + "validate_conversation_item_audio_formats", + "validate_input_sample_rate_hz", + "validate_realtime_item_truncate", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", + "validate_realtime_video_frames", + "validate_session_payload", + "wav_payload_to_pcm16", + "wire_value", +] diff --git a/vllm_omni/protocol/realtime/audio.py b/vllm_omni/protocol/realtime/audio.py new file mode 100644 index 000000000..5f163826a --- /dev/null +++ b/vllm_omni/protocol/realtime/audio.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import base64 +import binascii +import io +import math +import wave + +import numpy as np + +from vllm_omni.utils.audio_resample import StreamingAudioResampler + +try: + from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin +except ImportError: # pragma: no cover - audioop is removed in newer Python. + alaw2lin = lin2alaw = lin2ulaw = ulaw2lin = None + + +MIN_INPUT_SAMPLE_RATE_HZ = 8_000 +MAX_INPUT_SAMPLE_RATE_HZ = 192_000 + + +def validate_input_sample_rate_hz(sample_rate_hz: object) -> int: + if ( + isinstance(sample_rate_hz, bool) + or not isinstance(sample_rate_hz, int | float) + or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz)) + ): + raise ValueError("sample_rate_hz must be a finite integer") + rate = int(sample_rate_hz) + if rate != sample_rate_hz: + raise ValueError("sample_rate_hz must be an integer") + if not MIN_INPUT_SAMPLE_RATE_HZ <= rate <= MAX_INPUT_SAMPLE_RATE_HZ: + raise ValueError(f"sample_rate_hz must be between {MIN_INPUT_SAMPLE_RATE_HZ} and {MAX_INPUT_SAMPLE_RATE_HZ}") + return rate + + +def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes: + if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz: + return raw + samples = np.frombuffer(raw, dtype=" StreamingAudioResampler.max_polyphase_factor: + # The FIR size scales with the reduced ratio. Keep unusual client + # rates bounded while common audio rates use the high-quality path. + if samples.size <= 1: + return raw + target_size = max(1, int(round(samples.size * target_rate_hz / source_rate_hz))) + source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True) + target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True) + resampled = np.interp(target_x, source_x, samples) + return np.clip(resampled, -32_768, 32_767).astype(" bytes: + if ulaw2lin is not None: + return ulaw2lin(raw, 2) + data = np.frombuffer(raw, dtype=np.uint8) + value = np.bitwise_not(data).astype(np.int16) + sign = value & 0x80 + exponent = (value >> 4) & 0x07 + mantissa = value & 0x0F + sample = ((mantissa << 3) + 0x84) << exponent + return np.where(sign != 0, -(sample - 0x84), sample - 0x84).astype(" bytes: + if alaw2lin is not None: + return alaw2lin(raw, 2) + data = np.bitwise_xor(np.frombuffer(raw, dtype=np.uint8), 0x55).astype(np.int16) + sign = data & 0x80 + exponent = (data >> 4) & 0x07 + mantissa = data & 0x0F + sample = np.where(exponent == 0, (mantissa << 4) + 8, ((mantissa << 4) + 0x108) << (exponent - 1)) + return np.where(sign != 0, sample, -sample).astype(" bytes: + if lin2ulaw is not None: + return lin2ulaw(raw, 2) + pcm = np.clip(np.frombuffer(raw, dtype=" (0xFF << exp), exp + 1, exponent) + mantissa = (magnitude >> (exponent + 3)) & 0x0F + return (np.bitwise_not(sign | (exponent << 4) | mantissa) & 0xFF).astype(np.uint8).tobytes() + + +def encode_g711_alaw(raw: bytes) -> bytes: + if lin2alaw is not None: + return lin2alaw(raw, 2) + pcm = np.frombuffer(raw, dtype="= 0, 0x80, 0x00) + magnitude = np.abs(pcm) + exponent = np.zeros_like(magnitude) + for exp in range(1, 8): + exponent = np.where(magnitude >= (1 << (exp + 7)), exp, exponent) + mantissa = np.where( + exponent == 0, + (magnitude >> 4) & 0x0F, + (magnitude >> (exponent + 3)) & 0x0F, + ) + return ((sign | (exponent << 4) | mantissa) ^ 0x55).astype(np.uint8).tobytes() + + +def wav_payload_to_pcm16(raw: bytes) -> tuple[bytes | None, int | None]: + try: + with wave.open(io.BytesIO(raw), "rb") as wav_file: + channels = wav_file.getnchannels() + sample_width = wav_file.getsampwidth() + sample_rate_hz = wav_file.getframerate() + frames = wav_file.readframes(wav_file.getnframes()) + if sample_width != 2: + return None, sample_rate_hz + if channels <= 1: + return frames, sample_rate_hz + pcm = np.frombuffer(frames, dtype=" str: + """Package normalized mono float32 samples as a PCM16 WAV data payload.""" + rate = validate_input_sample_rate_hz(sample_rate_hz) + normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) + pcm16 = np.clip(np.rint(normalized * 32768.0), -32768, 32767).astype(" tuple[object, object, int | float | None]: + if not isinstance(audio, str) or not isinstance(fmt, str): + return audio, fmt, sample_rate_hz + normalized = fmt.lower() + if normalized not in {"pcm16", "pcm_s16le", "s16le", "g711_ulaw", "g711_alaw"}: + return audio, fmt, sample_rate_hz + if isinstance(sample_rate_hz, int | float): + sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) + try: + raw = base64.b64decode(audio.strip(), validate=False) + except (binascii.Error, ValueError): + return audio, fmt, sample_rate_hz + if normalized == "g711_ulaw": + raw = decode_g711_ulaw(raw) + sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 + elif normalized == "g711_alaw": + raw = decode_g711_alaw(raw) + sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 + elif len(raw) % 2: + return audio, fmt, sample_rate_hz + if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) != target_sample_rate_hz: + raw = resample_pcm16_mono( + raw, + source_rate_hz=int(sample_rate_hz), + target_rate_hz=target_sample_rate_hz, + ) + sample_rate_hz = target_sample_rate_hz + pcm = np.frombuffer(raw, dtype=" tuple[str, str, int | None]: + target = target_fmt.lower() + if target not in {"g711_ulaw", "g711_alaw"}: + return audio, source_fmt, source_sample_rate_hz + try: + raw = base64.b64decode(audio, validate=False) + except (binascii.Error, ValueError): + return audio, source_fmt, source_sample_rate_hz + source = source_fmt.lower() + if source == "wav": + pcm_raw, wav_rate = wav_payload_to_pcm16(raw) + if pcm_raw is None: + return audio, source_fmt, source_sample_rate_hz + raw = pcm_raw + source_sample_rate_hz = source_sample_rate_hz or wav_rate + elif source not in {"pcm", "pcm16", "pcm_s16le", "s16le"} or len(raw) % 2: + return audio, source_fmt, source_sample_rate_hz + target_rate = target_sample_rate_hz or 8_000 + if source_sample_rate_hz is not None: + raw = resample_pcm16_mono(raw, source_rate_hz=source_sample_rate_hz, target_rate_hz=target_rate) + encoded = encode_g711_ulaw(raw) if target == "g711_ulaw" else encode_g711_alaw(raw) + return base64.b64encode(encoded).decode("ascii"), target, target_rate diff --git a/vllm_omni/protocol/realtime/audio_input.py b/vllm_omni/protocol/realtime/audio_input.py new file mode 100644 index 000000000..634d32c34 --- /dev/null +++ b/vllm_omni/protocol/realtime/audio_input.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Decoding one ``input_audio_buffer.append``. + +An append is the hottest client event on a Realtime connection and the one with +the most wire surface: base64 audio in any accepted format, an optional rate, +optional camera frames, and a set of optional hints (``duration_ms``, +``is_speech``, a VAD probability, a transcript) that clients attach and servers +are free to use or ignore. + +:func:`decode_audio_append` turns that into :class:`RealtimeAudioAppend`, a +value object: bytes at 16 kHz ``pcm_f32le``, the surviving hints, and a speech +verdict. It is a pure function of the event plus the session's +:class:`~vllm_omni.protocol.realtime.session.RealtimeInputDefaults`, so any +consumer decodes an append the same way. Building a consumer's own command +object out of it is one constructor call --- see +``vllm_omni.engine.duplex.realtime_commands.build_append_audio``, which wraps it +for the duplex ``AppendAudio``. + +The speech verdict here is the *client-declared* one (explicit flags, a VAD +probability the client sent, or an RMS floor as a last resort). It is not +server-side VAD; that is a runtime concern and lives with whoever owns the +session. +""" + +from __future__ import annotations + +import base64 +import binascii +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import cast + +import numpy as np + +from vllm_omni.protocol.realtime.audio import convert_input_audio_with_rate +from vllm_omni.protocol.realtime.errors import RealtimeProtocolError +from vllm_omni.protocol.realtime.formats import ( + is_supported_realtime_input_format, + parse_realtime_audio_format, +) +from vllm_omni.protocol.realtime.items import validate_realtime_video_frames +from vllm_omni.protocol.realtime.session import RealtimeInputDefaults + +__all__ = [ + "REALTIME_INPUT_HINT_KEYS", + "RealtimeAudioAppend", + "copy_realtime_input_hints", + "decode_audio_append", + "input_explicitly_non_speech", + "input_looks_like_speech", +] + +#: Wire hint keys copied verbatim from a client append onto the internal payload. +REALTIME_INPUT_HINT_KEYS = ( + "duration_ms", + "audio_duration_ms", + "audio_start_ms", + "audio_end_ms", + "is_speech", + "speech", + "speech_probability", + "vad", + "overlap_action", + "overlap", + "force_barge_in", + "force_listen", + "text", + "transcript", +) + + +@dataclass(frozen=True, slots=True) +class RealtimeAudioAppend: + """One decoded ``input_audio_buffer.append``. + + ``audio`` is raw bytes in ``format`` at ``sample_rate_hz`` --- after + :func:`decode_audio_append` that is 16 kHz ``pcm_f32le`` for every input + format the codec converts. + """ + + audio: bytes + format: str + sample_rate_hz: int | None = None + is_speech: bool | None = None + video_frames: tuple[str, ...] = () + duration_ms: int | None = None + audio_end_ms: int | None = None + #: Model-neutral hints carried through from the wire (rms, vad, transcript hints ...). + hints: dict[str, object] = field(default_factory=dict) + #: Client correlation id (OpenAI ``event_id``), echoed on error events. + event_id: str | None = None + + +# ---- speech hints ---- + + +def input_explicitly_non_speech(event: Mapping[str, object]) -> bool: + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return not value + vad = event.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return not value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) < 0.5 + probability = event.get("speech_probability") + return isinstance(probability, int | float) and float(probability) < 0.5 + + +def input_looks_like_speech( + event: Mapping[str, object], + *, + audio: object, + fmt: object, + overlap_silence_rms: float, +) -> bool: + if input_explicitly_non_speech(event): + return False + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return value + vad = event.get("vad") + if isinstance(vad, dict): + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = event.get("speech_probability") + if isinstance(probability, int | float): + return float(probability) >= 0.5 + if fmt != "pcm_f32le": + return True + if isinstance(audio, bytes | bytearray): + raw = bytes(audio) + elif isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return True + else: + return True + if len(raw) < 4 or len(raw) % 4 != 0: + return True + samples = np.frombuffer(raw, dtype=np.float32) + if samples.size == 0: + return False + rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) + threshold = event.get("overlap_silence_rms") + if not isinstance(threshold, int | float): + vad = event.get("vad") + if isinstance(vad, dict): + threshold = vad.get("silence_rms") + silence_rms = float(threshold) if isinstance(threshold, int | float) else overlap_silence_rms + return rms >= max(0.0, silence_rms) + + +def copy_realtime_input_hints(source: Mapping[str, object], target: dict[str, object]) -> None: + for key in REALTIME_INPUT_HINT_KEYS: + if key in source: + target[key] = source[key] + + +# ---- append decoding ---- + + +def decode_audio_append( + event: Mapping[str, object], + *, + defaults: RealtimeInputDefaults, + hints_source: Mapping[str, object] | None = None, +) -> RealtimeAudioAppend: + """Validate, convert (to 16 kHz ``pcm_f32le``) and pack one audio append. + + ``hints_source`` is the enclosing payload when the audio arrives inside + something larger than a bare append --- a ``conversation.item.create`` + audio part, say --- so its hints apply unless the part overrides them. + + Raises :class:`RealtimeProtocolError` for an unsupported format, undecodable + audio or invalid camera frames. + """ + event_id = cast("str", event.get("event_id")) if isinstance(event.get("event_id"), str) else None + audio = event.get("audio") or event.get("delta") + fmt, format_rate = parse_realtime_audio_format( + event.get("format") or event.get("input_audio_format") or defaults.input_audio_format + ) + sample_rate_hz = ( + event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or defaults.input_sample_rate_hz + ) + if not is_supported_realtime_input_format(fmt): + raise RealtimeProtocolError( + f"Unsupported input_audio_format: {fmt}", + code="unsupported_audio_format", + event_id=event_id, + ) + try: + audio, fmt, sample_rate_hz = convert_input_audio_with_rate( + audio, + fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, + ) + except ValueError as exc: + raise RealtimeProtocolError(str(exc), code="bad_event", event_id=event_id) from exc + hints: dict[str, object] = {} + if hints_source is not None: + copy_realtime_input_hints(hints_source, hints) + copy_realtime_input_hints(event, hints) + looks_like_speech = input_looks_like_speech( + {**(hints_source or {}), **event}, + audio=audio, + fmt=fmt, + overlap_silence_rms=defaults.overlap_silence_rms, + ) + video_frames: tuple[str, ...] = () + raw_frames = event.get("video_frames") + if raw_frames is not None: + frames_error = validate_realtime_video_frames(raw_frames, event.get("max_slice_nums")) + if frames_error is not None: + raise RealtimeProtocolError(frames_error, code="invalid_video_frames", event_id=event_id) + if isinstance(raw_frames, list): + video_frames = tuple(frame for frame in raw_frames if isinstance(frame, str) and frame) + duration_ms = hints.get("duration_ms", hints.get("audio_duration_ms")) + audio_end_ms = hints.get("audio_end_ms") + try: + audio_bytes = base64.b64decode(audio, validate=True) if isinstance(audio, str) and audio else b"" + except (binascii.Error, ValueError) as exc: + raise RealtimeProtocolError("input audio is not valid base64", code="bad_audio", event_id=event_id) from exc + return RealtimeAudioAppend( + event_id=event_id, + audio=audio_bytes, + format=str(fmt), + sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, + is_speech=looks_like_speech, + video_frames=video_frames, + duration_ms=int(duration_ms) if isinstance(duration_ms, int | float) else None, + audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else None, + hints=hints, + ) diff --git a/vllm_omni/protocol/realtime/capabilities.py b/vllm_omni/protocol/realtime/capabilities.py new file mode 100644 index 000000000..658a57e56 --- /dev/null +++ b/vllm_omni/protocol/realtime/capabilities.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""What a Realtime consumer can serve, and the session check that reads it. + +``session.update`` is the one client event whose acceptance depends on who is +behind the socket. The audio formats a server can decode, and whether it can do +server-side turn detection at all, are properties of the *consumer* --- the +duplex engine, or some other runtime --- not of the protocol. The codec still +has to do the checking, because it owns the error envelope and the order the +checks run in. + +:class:`RealtimeProtocolCapabilities` is how a consumer states those answers. +It is the post-#7413 form of the ``supports(turn_detection, formats)`` half of +RFC #6592's proposed ``RealtimeModelAdapter``. The other three halves of that +proposal --- building a prompt, starting a response, cancelling one --- are not +here and must not come here: for a duplex model they are already owned by +``DuplexModelPlugin`` together with ``DuplexSessionRunner``, and a second +interface over the same responsibility is exactly the fork this layer exists to +prevent. + +The duplex binding of this object is +``vllm_omni.engine.duplex.realtime_commands.DUPLEX_REALTIME_CAPABILITIES``. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field + +from vllm_omni.protocol.realtime.formats import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_OUTPUT_AUDIO_FORMATS, + validate_realtime_session_audio_formats, +) + +__all__ = [ + "RealtimeProtocolCapabilities", + "RealtimeSessionRejection", + "validate_session_payload", +] + +#: A consumer's answer to "can you do this turn detection?": an error message, or None. +TurnDetectionValidator = Callable[[Mapping[str, object]], str | None] + + +@dataclass(frozen=True, slots=True) +class RealtimeProtocolCapabilities: + """One consumer's answer to what a ``session`` object may ask for. + + The defaults are permissive, not restrictive: every format the codec can + decode, and **no turn-detection validation at all**. A consumer that leaves + ``validate_turn_detection`` at ``None`` therefore accepts whatever + ``turn_detection`` the session object carries --- including values it cannot + actually serve, such as ``semantic_vad``. A consumer that supports only some + turn-detection modes (or none) must supply its own validator; see + ``vllm_omni.engine.duplex.realtime_commands.DUPLEX_REALTIME_CAPABILITIES``. + """ + + input_audio_formats: frozenset[str] = field(default_factory=lambda: frozenset(REALTIME_INPUT_AUDIO_FORMATS)) + output_audio_formats: frozenset[str] = field(default_factory=lambda: frozenset(REALTIME_OUTPUT_AUDIO_FORMATS)) + #: Consumer-specific ``turn_detection`` validation. ``None`` means the + #: consumer takes whatever the session object says without a check of its own. + validate_turn_detection: TurnDetectionValidator | None = None + + +@dataclass(frozen=True, slots=True) +class RealtimeSessionRejection: + """Why a session object was refused: the internal error code and the message. + + ``param`` names the offending session field where there is one, for the + ``error.param`` slot of the OpenAI error envelope. + """ + + code: str + message: str + param: str | None = None + + +def validate_session_payload( + session_payload: Mapping[str, object], + *, + capabilities: RealtimeProtocolCapabilities, +) -> RealtimeSessionRejection | None: + """Check a ``session`` object against one consumer; ``None`` when it is acceptable. + + Audio formats are checked before turn detection, so a session object that + is wrong in both ways reports the format problem --- keep that order, it is + what clients already see. + """ + format_error = validate_realtime_session_audio_formats( + session_payload, + input_audio_formats=capabilities.input_audio_formats, + output_audio_formats=capabilities.output_audio_formats, + ) + if format_error is not None: + return RealtimeSessionRejection(code="unsupported_audio_format", message=format_error) + validate_turn_detection = capabilities.validate_turn_detection + if validate_turn_detection is not None: + turn_detection_error = validate_turn_detection(session_payload) + if turn_detection_error is not None: + return RealtimeSessionRejection( + code="unsupported_turn_detection", message=turn_detection_error, param="turn_detection" + ) + return None diff --git a/vllm_omni/protocol/realtime/commands.py b/vllm_omni/protocol/realtime/commands.py new file mode 100644 index 000000000..590f3ee93 --- /dev/null +++ b/vllm_omni/protocol/realtime/commands.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime client events as typed commands. + +A Realtime server's first job is turning a client JSON event into something +typed. These are the 10 client events OpenAI's Realtime API defines, as frozen +dataclasses carrying the fields that survive decoding. + +vLLM-Omni's own client events (barge-in, playback acknowledgement, text append, +heartbeat, explicit close, generic turn signals) are a separate vocabulary in +``vllm_omni.protocol.duplex.commands``. + +Decoding a payload *into* these stays with the consumer, because which command +a given event becomes can depend on what that consumer supports --- see +``vllm_omni.engine.duplex.realtime_commands.translate_realtime_command``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import ClassVar + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RealtimeCommand: + """One decoded OpenAI Realtime client event. + + ``wire_type`` is the event type the client sent; the typed fields are what + survived decoding, so a consumer reads them instead of re-parsing JSON. + + There is deliberately no rendering method here. A server receives commands, + it does not emit them, and how a runtime represents one internally is its + own business --- the duplex engine renders its mailbox dictionary in + ``vllm_omni.engine.duplex.commands``. + """ + + #: OpenAI Realtime client event this command decodes from. + wire_type: ClassVar[str] = "" + #: Client correlation id (OpenAI ``event_id``), echoed on error events. + event_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class UpdateSession(RealtimeCommand): + """``session.update``: a Realtime ``session`` object patch.""" + + wire_type: ClassVar[str] = "session.update" + patch: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendAudio(RealtimeCommand): + wire_type: ClassVar[str] = "input_audio_buffer.append" + #: Raw audio bytes in ``format`` at ``sample_rate_hz`` (base64 only on the wire). + audio: bytes + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Commit(RealtimeCommand): + wire_type: ClassVar[str] = "input_audio_buffer.commit" + final: bool = True + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearInput(RealtimeCommand): + wire_type: ClassVar[str] = "input_audio_buffer.clear" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearOutputAudio(RealtimeCommand): + wire_type: ClassVar[str] = "output_audio_buffer.clear" + #: Explicit response to clear; ``None`` targets the active/last response. + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateResponse(RealtimeCommand): + wire_type: ClassVar[str] = "response.create" + #: Raw Realtime ``response`` object (instructions, voice, modalities, ...). + options: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelResponse(RealtimeCommand): + wire_type: ClassVar[str] = "response.cancel" + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateItem(RealtimeCommand): + """``conversation.item.create`` (history injection or function-call output).""" + + wire_type: ClassVar[str] = "conversation.item.create" + item: Mapping[str, object] + previous_item_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DeleteItem(RealtimeCommand): + wire_type: ClassVar[str] = "conversation.item.delete" + item_id: str + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TruncateItem(RealtimeCommand): + wire_type: ClassVar[str] = "conversation.item.truncate" + item_id: str + audio_end_ms: int + content_index: int = 0 + + +__all__ = [ + "AppendAudio", + "CancelResponse", + "ClearInput", + "ClearOutputAudio", + "Commit", + "CreateItem", + "CreateResponse", + "DeleteItem", + "RealtimeCommand", + "TruncateItem", + "UpdateSession", +] diff --git a/vllm_omni/protocol/realtime/errors.py b/vllm_omni/protocol/realtime/errors.py new file mode 100644 index 000000000..f9ba67bf4 --- /dev/null +++ b/vllm_omni/protocol/realtime/errors.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The OpenAI Realtime error envelope. + +A Realtime server reports a rejected client event as one ``error`` event whose +``error`` object carries ``type`` / ``code`` / ``message`` (plus the client's +``event_id`` and an optional ``param``). ``type`` is not free-form: it is one +of OpenAI's four buckets, derived from our internal code by +:data:`REALTIME_ERROR_TYPES_BY_CODE`. + +OpenAI standardises the three ``type`` classes but not the *codes*, so the code +vocabulary is a consumer's own. vLLM-Omni's lives in +``vllm_omni.protocol.duplex.errors`` (Tier 3 in +``docs/serving/realtime_duplex_api.md``); what stays here is the exception type +every consumer raises, so the envelope is rendered identically whoever built it. +""" + +from __future__ import annotations + +__all__ = ["RealtimeProtocolError"] + + +class RealtimeProtocolError(ValueError): + """A client payload could not be decoded into a valid Realtime intent. + + ``code`` is the internal error code that :data:`REALTIME_ERROR_TYPES_BY_CODE` + maps to an OpenAI ``error.type``; ``event_id`` is the *client* event id the + error answers. ``vllm_omni.engine.duplex.commands.DuplexCommandError`` is + the duplex specialization, so a consumer catching either sees the same + three attributes. + """ + + def __init__(self, message: str, *, code: str = "bad_event", event_id: str | None = None) -> None: + super().__init__(message) + self.code = code + self.event_id = event_id diff --git a/vllm_omni/protocol/realtime/events.py b/vllm_omni/protocol/realtime/events.py new file mode 100644 index 000000000..47bab6372 --- /dev/null +++ b/vllm_omni/protocol/realtime/events.py @@ -0,0 +1,467 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime server events: typed classes and their wire rendering. + +Every event is a frozen dataclass with explicit fields, and the Realtime wire +JSON is *derived* from those fields by :meth:`RealtimeEvent.to_realtime`. That +rendering is pure --- it never consults session state --- which is why it can +live here, away from whoever owns the session. The stateful part (response and +item ids, content-part bookkeeping) is consumed when the events are +*constructed*, by the consumer. + +This module holds the 30 events OpenAI's Realtime API defines. vLLM-Omni's own +additions to the wire (listen/speak turn-taking, playback acknowledgement, +session resume) are a separate vocabulary in +``vllm_omni.protocol.duplex.events``, so a consumer that only speaks GA does +not inherit them. + +``session_id`` / ``epoch`` are bound by the consumer when an event leaves it; +producers build events with the defaults. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import ClassVar +from uuid import uuid4 + + +def new_event_id() -> str: + return f"event_{uuid4().hex}" + + +_IDENTITY_FIELDS = frozenset({"session_id", "epoch", "event_id"}) + + +def wire_value(value: object) -> object: + if isinstance(value, Mapping): + return {str(k): wire_value(v) for k, v in value.items()} + if isinstance(value, tuple | list): + return [wire_value(v) for v in value] + return value + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RealtimeEvent: + """Base of every public session event.""" + + #: Realtime ``type`` this class renders to. + wire_type: ClassVar[str] = "" + #: Fields that are omitted from the wire object when ``None``. + optional_wire_fields: ClassVar[frozenset[str]] = frozenset() + + session_id: str = "" + epoch: int | None = None + #: Server event id (OpenAI ``event_id``); generated at construction. + event_id: str = field(default_factory=new_event_id) + + # ---- identity ---- + + @property + def type(self) -> str: + return self.wire_type + + @property + def is_terminal(self) -> bool: + return False + + # ---- generic accessors (None when the class has no such field) ---- + + @property + def response_id(self) -> str | None: + return None + + @property + def item_id(self) -> str | None: + return None + + @property + def text(self) -> str | None: + return None + + @property + def audio(self) -> bytes | None: + return None + + # ---- wire rendering (pure) ---- + + def _wire_fields(self) -> dict[str, object]: + data: dict[str, object] = {} + for f in fields(self): + if f.name in _IDENTITY_FIELDS: + continue + value = getattr(self, f.name) + if value is None and f.name in self.optional_wire_fields: + continue + data[f.name] = wire_value(value) + return data + + def to_realtime(self) -> dict[str, object]: + """The Realtime wire JSON object for this event (derived from the fields).""" + return {"type": self.wire_type, "event_id": self.event_id, **self._wire_fields()} + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"{type(self).__name__}(type={self.wire_type!r}, session_id={self.session_id!r}, epoch={self.epoch!r})" + + +# ---- shared bases ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseScopedEvent(RealtimeEvent): + """Events addressed to one response (and usually one output item / content part).""" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + content_index: int = 0 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ItemEvent(RealtimeEvent): + previous_item_id: str | None = None + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +# ---- GA server events ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionCreated(RealtimeEvent): + wire_type = "session.created" + + session: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionUpdated(RealtimeEvent): + wire_type = "session.updated" + + session: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseCreated(RealtimeEvent): + wire_type = "response.created" + + response: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemAdded(RealtimeEvent): + wire_type = "response.output_item.added" + + response_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemDone(OutputItemAdded): + wire_type = "response.output_item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartAdded(ResponseScopedEvent): + wire_type = "response.content_part.added" + + part: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartDone(ContentPartAdded): + wire_type = "response.content_part.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDelta(ResponseScopedEvent): + wire_type = "response.output_audio.delta" + + #: Base64 audio in ``format``. + delta: str = "" + + @property + def audio(self) -> bytes | None: + try: + return base64.b64decode(self.delta) + except (ValueError, TypeError): + return None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDone(ResponseScopedEvent): + wire_type = "response.output_audio.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDelta(ResponseScopedEvent): + wire_type = "response.output_audio_transcript.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDone(ResponseScopedEvent): + wire_type = "response.output_audio_transcript.done" + + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDelta(ResponseScopedEvent): + wire_type = "response.output_text.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDone(ResponseScopedEvent): + wire_type = "response.output_text.done" + + text: str = "" # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDelta(RealtimeEvent): + wire_type = "response.function_call_arguments.delta" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + delta: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDone(RealtimeEvent): + wire_type = "response.function_call_arguments.done" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + arguments: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseDone(RealtimeEvent): + wire_type = "response.done" + + response: Mapping[str, object] = field(default_factory=dict) + + @property + def status(self) -> str | None: + value = self.response.get("status") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RateLimitsUpdated(RealtimeEvent): + wire_type = "rate_limits.updated" + + rate_limits: tuple[Mapping[str, object], ...] = () + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCommitted(RealtimeEvent): + wire_type = "input_audio_buffer.committed" + + previous_item_id: str | None = None + item_id: str | None = None # type: ignore[assignment] + + def _wire_fields(self) -> dict[str, object]: + return {"previous_item_id": self.previous_item_id, "item_id": self.item_id} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCleared(RealtimeEvent): + wire_type = "input_audio_buffer.cleared" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStarted(RealtimeEvent): + wire_type = "input_audio_buffer.speech_started" + + audio_start_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStopped(RealtimeEvent): + wire_type = "input_audio_buffer.speech_stopped" + + audio_end_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemAdded(_ItemEvent): + wire_type = "conversation.item.added" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemCreated(_ItemEvent): + wire_type = "conversation.item.created" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDone(_ItemEvent): + wire_type = "conversation.item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemRetrieved(RealtimeEvent): + wire_type = "conversation.item.retrieved" + + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDeleted(RealtimeEvent): + wire_type = "conversation.item.deleted" + + item_id: str | None = None # type: ignore[assignment] + + def _wire_fields(self) -> dict[str, object]: + return {"item_id": self.item_id} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemTruncated(RealtimeEvent): + wire_type = "conversation.item.truncated" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + audio_end_ms: int = 0 + + def _wire_fields(self) -> dict[str, object]: + return { + "item_id": self.item_id, + "content_index": self.content_index, + "audio_end_ms": self.audio_end_ms, + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputTranscriptionCompleted(RealtimeEvent): + wire_type = "conversation.item.input_audio_transcription.completed" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputAudioCleared(RealtimeEvent): + wire_type = "output_audio_buffer.cleared" + + response_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ErrorEvent(RealtimeEvent): + wire_type = "error" + + code: str = "internal_error" + message: str = "" + #: Client ``event_id`` this error answers (OpenAI ``error.event_id``). + related_event_id: str | None = None + param: str | None = None + + @property + def error_type(self) -> str: + """OpenAI's error class for this error. + + OpenAI standardises only the three classes, not the codes, so the base + reports the default. A consumer with its own code vocabulary overrides + this --- see ``vllm_omni.protocol.duplex.events.ErrorEvent``. + """ + return "invalid_request_error" + + @property + def error(self) -> dict[str, object]: + error: dict[str, object] = {"type": self.error_type, "code": self.code, "message": self.message} + if self.related_event_id: + error["event_id"] = self.related_event_id + if self.param: + error["param"] = self.param + return error + + @property + def text(self) -> str | None: + return self.message + + def _wire_fields(self) -> dict[str, object]: + return {"error": self.error} + + +__all__ = [ + "AudioDelta", + "AudioDone", + "ContentPartAdded", + "ContentPartDone", + "ErrorEvent", + "FunctionCallArgumentsDelta", + "FunctionCallArgumentsDone", + "InputCleared", + "InputCommitted", + "InputTranscriptionCompleted", + "ItemAdded", + "ItemCreated", + "ItemDeleted", + "ItemDone", + "ItemRetrieved", + "ItemTruncated", + "OutputAudioCleared", + "OutputItemAdded", + "OutputItemDone", + "RateLimitsUpdated", + "RealtimeEvent", + "ResponseCreated", + "ResponseDone", + "ResponseScopedEvent", + "SessionCreated", + "SessionUpdated", + "SpeechStarted", + "SpeechStopped", + "TextDelta", + "TextDone", + "TranscriptDelta", + "TranscriptDone", + "new_event_id", + "wire_value", +] diff --git a/vllm_omni/protocol/realtime/formats.py b/vllm_omni/protocol/realtime/formats.py new file mode 100644 index 000000000..161aab11b --- /dev/null +++ b/vllm_omni/protocol/realtime/formats.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Realtime audio-format negotiation: what a client may declare and how it is read. + +A Realtime client declares audio formats in three places --- the session object +(``session.audio.input.format`` and the older flat ``input_audio_format``), a +``response.create`` override, and a ``conversation.item.create`` audio part --- +and may spell one format several ways (``audio/pcm``, ``pcm16``, ``s16le``). +This module is the single place that normalizes those spellings and says which +are supported. Folding a session's declaration into the per-connection append +defaults is the job of ``session.py``, which builds on this module. + +Everything here is a pure function of wire payloads: no session state, no model, +no engine. +""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping + +__all__ = [ + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "is_supported_realtime_input_format", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_output_format", + "validate_conversation_item_audio_formats", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", +] + +REALTIME_INPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "pcm_f32le", + "g711_ulaw", + "g711_alaw", +} +REALTIME_OUTPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "wav", + "pcm", + "g711_ulaw", + "g711_alaw", +} + + +def parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: + def normalize_format(fmt: str) -> str: + normalized = fmt.lower() + if normalized in {"audio/pcm", "pcm"}: + return "pcm16" + if normalized in {"audio/wav", "wav"}: + return "wav" + if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: + return "pcm16" + if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: + return "pcm_f32le" + if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: + return "g711_ulaw" + if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: + return "g711_alaw" + return fmt + + if isinstance(raw_format, str): + return normalize_format(raw_format), None + if not isinstance(raw_format, dict): + return raw_format, None + rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") + sample_rate_hz = int(rate) if isinstance(rate, int | float) and rate > 0 else None + fmt = raw_format.get("type") or raw_format.get("format") + if not isinstance(fmt, str): + return raw_format, sample_rate_hz + return normalize_format(fmt), sample_rate_hz + + +def realtime_output_format(duplex_format: object) -> str: + if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: + return duplex_format.lower() + if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": + return "pcm16" + return str(duplex_format or "wav") + + +def is_supported_realtime_input_format(fmt: object) -> bool: + return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS + + +def realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: + payload: dict[str, object] = {"type": "audio/pcm"} + elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": + payload = {"type": "audio/pcm_f32le"} + elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": + payload = {"type": "audio/g711_ulaw"} + elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": + payload = {"type": "audio/g711_alaw"} + else: + payload = {"type": "audio/wav"} + if sample_rate_hz is not None: + payload["rate"] = int(sample_rate_hz) + return payload + + +def validate_realtime_session_audio_formats( + session_payload: Mapping[str, object], + *, + input_audio_formats: Collection[str] | None = None, + output_audio_formats: Collection[str] | None = None, +) -> str | None: + """Reject a session object that declares an audio format we cannot serve. + + The format sets default to everything the codec can decode; a consumer that + serves a narrower set passes its own (see + ``vllm_omni.protocol.realtime.capabilities``). + """ + supported_input = REALTIME_INPUT_AUDIO_FORMATS if input_audio_formats is None else input_audio_formats + supported_output = REALTIME_OUTPUT_AUDIO_FORMATS if output_audio_formats is None else output_audio_formats + audio_config = session_payload.get("audio") + input_format: object = session_payload.get("input_audio_format") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + parsed_input, _ = parse_realtime_audio_format(input_format) + if input_format is not None and not (isinstance(parsed_input, str) and parsed_input.lower() in supported_input): + return f"Unsupported input_audio_format: {input_format}" + + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not (isinstance(parsed_output, str) and parsed_output.lower() in supported_output): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_realtime_response_audio_formats(response_payload: Mapping[str, object]) -> str | None: + output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") + audio_config = response_payload.get("audio") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not ( + isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_conversation_item_audio_formats(item: object) -> str | None: + if not isinstance(item, dict): + return None + content = item.get("content") + if not isinstance(content, list): + return None + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in {"input_audio", "audio"}: + continue + raw_format = part.get("format") + parsed_format, _ = parse_realtime_audio_format(raw_format) + if raw_format is not None and not ( + isinstance(parsed_format, str) and parsed_format.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + return f"Unsupported input_audio format in conversation.item.create: {raw_format}" + return None diff --git a/vllm_omni/protocol/realtime/items.py b/vllm_omni/protocol/realtime/items.py new file mode 100644 index 000000000..cbcb9c652 --- /dev/null +++ b/vllm_omni/protocol/realtime/items.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Realtime conversation items: shape, truncation and the transcript inside one. + +A Realtime conversation is a list of *items* --- a user message, an assistant +message, a function call --- each with a ``content`` list of typed parts. The +client creates them (``conversation.item.create``), truncates an assistant +item's audio after a barge-in (``conversation.item.truncate``), and the server +echoes them back inside ``conversation.item.*`` and ``response.done`` events. + +The rules in this module are the ones both sides have to agree on: which +defaults a bare item gets, how far a transcript is cut when the client says it +only heard N milliseconds of audio, and what makes a truncate request invalid. +Item *storage* --- which items a session currently holds, in what order --- is +the consumer's state, not this module's. + +``validate_realtime_video_frames`` sits here too: camera frames ride along an +``input_audio_buffer.append`` in the vLLM-Omni extension of the protocol, and +validating them is the same kind of wire-shape check. +""" + +from __future__ import annotations + +import base64 +import binascii +from collections.abc import Mapping +from uuid import uuid4 + +__all__ = [ + "input_transcript_from_item", + "normalize_conversation_item", + "text_chars_for_audio_ms_from_marks", + "truncate_realtime_item_content", + "validate_realtime_item_truncate", + "validate_realtime_video_frames", +] + + +def validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: + """Validate omni-duplex camera frames on input_audio_buffer.append. + + Wire contract matches the official MiniCPM-o duplex loop: one base + base64 JPEG per ~1 s audio chunk, optionally followed by that unit's + stacked composite tiling the sub-frames captured inside it (at most 2 + images either way). A caller-supplied ``max_slice_nums`` is rejected rather + than silently ignored: slicing is Stage 0's decision here, and Stage 0 + already applies the official HD suggestion for a stacked unit + (``max_slice_nums=[2, 1]``). The wire simply does not let the client choose + it. + """ + if max_slice_nums not in (None, 1): + return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" + if not isinstance(video_frames, list): + return "video_frames must be a list of base64-encoded images" + frames = [frame for frame in video_frames if frame is not None] + if len(frames) > 2: + return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" + for frame in frames: + if not isinstance(frame, str) or not frame: + return "video_frames entries must be non-empty base64 strings" + if len(frame) > 4_000_000: + return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" + try: + header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) + except (binascii.Error, ValueError): + return "video_frames entries must be valid base64" + if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): + return "video_frames entries must be JPEG or PNG images" + return None + + +def normalize_conversation_item(item: Mapping[str, object]) -> dict[str, object]: + normalized = dict(item) + normalized.setdefault("id", f"item_{uuid4().hex}") + normalized.setdefault("object", "realtime.item") + normalized.setdefault("type", "message") + normalized.setdefault("status", "completed") + if "role" not in normalized and normalized.get("type") == "message": + normalized["role"] = "user" + if not isinstance(normalized.get("content"), list): + normalized["content"] = [] + return normalized + + +def text_chars_for_audio_ms_from_marks( + audio_end_ms: int, + text_len: int, + marks: list[object], + *, + final_ms: object | None = None, +) -> int: + if text_len <= 0: + return 0 + clean_marks: list[tuple[int, int]] = [] + for mark in marks: + if not isinstance(mark, dict): + continue + raw_text_chars = mark.get("text_chars") + raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) + if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): + continue + clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) + if not clean_marks: + return 0 if audio_end_ms <= 0 else text_len + clean_marks.sort(key=lambda item: item[0]) + audio_end_ms = max(0, int(audio_end_ms)) + if audio_end_ms <= 0: + return 0 + previous_ms = 0 + previous_chars = 0 + for mark_ms, mark_chars in clean_marks: + mark_ms = max(previous_ms, mark_ms) + mark_chars = max(previous_chars, min(text_len, mark_chars)) + if audio_end_ms <= mark_ms: + if mark_ms <= previous_ms: + return mark_chars + ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) + return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) + previous_ms = mark_ms + previous_chars = mark_chars + if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: + if audio_end_ms >= int(final_ms): + return text_len + ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) + return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) + return text_len if audio_end_ms >= previous_ms else previous_chars + + +def truncate_realtime_item_content(item: dict[str, object], *, content_index: int, audio_end_ms: int) -> None: + content = item.get("content") + if not isinstance(content, list) or not content: + return + index = max(0, int(content_index)) + if index >= len(content): + return + part = content[index] + if not isinstance(part, dict): + return + transcript = part.get("transcript") + if not isinstance(transcript, str) or not transcript: + return + marks = part.get("audio_text_marks") + if isinstance(marks, list): + keep_chars = text_chars_for_audio_ms_from_marks( + audio_end_ms, + len(transcript), + marks, + final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), + ) + part["transcript"] = transcript[:keep_chars].rstrip() + return + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and duration_ms > 0: + keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) + part["transcript"] = transcript[:keep_chars].rstrip() + elif audio_end_ms <= 0: + part["transcript"] = "" + + +def validate_realtime_item_truncate(item: Mapping[str, object], *, content_index: int, audio_end_ms: int) -> str | None: + if item.get("type") != "message" or item.get("role") != "assistant": + return "conversation.item.truncate only supports assistant message items" + if audio_end_ms < 0: + return "conversation.item.truncate requires non-negative audio_end_ms" + content = item.get("content") + if not isinstance(content, list) or not content: + return None + index = max(0, int(content_index)) + if index >= len(content): + return f"conversation.item.truncate content_index out of range: {content_index}" + part = content[index] + if not isinstance(part, dict): + return "conversation.item.truncate target content part is invalid" + if part.get("type") not in {"audio", "output_audio"}: + return "conversation.item.truncate target content part must be audio" + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): + return "conversation.item.truncate audio_end_ms exceeds item audio duration" + return None + + +def input_transcript_from_item(item: Mapping[str, object]) -> str: + content = item.get("content") + if not isinstance(content, list): + return "" + parts: list[str] = [] + for part in content: + if not isinstance(part, dict): + continue + for key in ("transcript", "text"): + value = part.get(key) + if isinstance(value, str) and value: + if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: + parts.append(value) + break + return "".join(parts).strip() diff --git a/vllm_omni/protocol/realtime/session.py b/vllm_omni/protocol/realtime/session.py new file mode 100644 index 000000000..1c73ea926 --- /dev/null +++ b/vllm_omni/protocol/realtime/session.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Reading the Realtime ``session`` object. + +``session.update`` carries one JSON object that a client may spell in either of +two shapes: the GA nest (``session.audio.input.format``, +``session.audio.output.sample_rate_hz``) or the older flat beta keys +(``input_audio_format``, ``sample_rate_hz``). Everything downstream --- the +append decoder, the projection that renders ``session.created`` back out, a +model's runtime configuration --- has to agree on what a given session object +means, so the readers live here once. + +:class:`RealtimeInputDefaults` is the resolved answer for the input path: the +format, sample rate and silence threshold an ``input_audio_buffer.append`` +inherits when it omits them. A connection keeps one of these and re-derives it +on every ``session.update``. + +``overlap_policy`` / ``playback_commit_policy`` are vLLM-Omni extension fields +on the session object rather than GA OpenAI fields. They are read here because +this module owns what the wire may say; what a server *does* with them is the +consumer's business. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Any + +from vllm_omni.protocol.realtime.formats import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_OUTPUT_AUDIO_FORMATS, + parse_realtime_audio_format, + realtime_output_format, +) + +__all__ = [ + "RealtimeInputDefaults", + "apply_realtime_session_defaults", + "input_audio_transcription_config", + "json_safe_realtime_payload", + "realtime_max_output_tokens", + "realtime_overlap_fields", +] + + +@dataclass(frozen=True, slots=True) +class RealtimeInputDefaults: + """Session-level wire defaults an append may omit (derived from the session object).""" + + input_audio_format: str = "pcm16" + input_sample_rate_hz: int = 16000 + output_audio_format: str = "pcm16" + output_sample_rate_hz: int | None = None + overlap_silence_rms: float = 0.003 + + def with_session_payload(self, session_payload: Mapping[str, object]) -> RealtimeInputDefaults: + """Return defaults updated from a Realtime ``session`` object (session.update).""" + values = apply_realtime_session_defaults(self, session_payload) + return values + + +def apply_realtime_session_defaults( + defaults: RealtimeInputDefaults, + session_payload: Mapping[str, object], +) -> RealtimeInputDefaults: + """Derive the wire defaults a Realtime ``session`` object declares.""" + input_format: object = session_payload.get("input_audio_format") + audio_config = session_payload.get("audio") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + input_format, input_rate = parse_realtime_audio_format(input_format) + # Any: heterogeneous field values splatted into ``dataclasses.replace`` below. + updates: dict[str, Any] = {} + if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: + updates["input_audio_format"] = input_format + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + output_rate_raw: object | None = None + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") + output_format, output_rate = parse_realtime_audio_format(output_format) + if output_rate is None and isinstance(output_rate_raw, int | float) and output_rate_raw > 0: + output_rate = int(output_rate_raw) + if isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + updates["output_audio_format"] = realtime_output_format(output_format) + sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") + if sample_rate is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") + if sample_rate is None: + sample_rate = input_rate + if isinstance(sample_rate, int | float) and sample_rate > 0: + updates["input_sample_rate_hz"] = int(sample_rate) + if isinstance(output_rate, int | float) and output_rate > 0: + updates["output_sample_rate_hz"] = int(output_rate) + overlap_fields = realtime_overlap_fields(session_payload) + overlap_silence_rms = overlap_fields.get("overlap_silence_rms") + if isinstance(overlap_silence_rms, int | float): + updates["overlap_silence_rms"] = max(0.0, float(overlap_silence_rms)) + return replace(defaults, **updates) if updates else defaults + + +def realtime_overlap_fields(session_payload: Mapping[str, object]) -> dict[str, object]: + fields: dict[str, object] = {} + if isinstance(session_payload.get("overlap_policy"), str): + fields["overlap_policy"] = session_payload["overlap_policy"] + for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): + value = session_payload.get(key) + if isinstance(value, int | float): + fields[key] = value + if isinstance(session_payload.get("playback_commit_policy"), str): + fields["playback_commit_policy"] = session_payload["playback_commit_policy"] + return fields + + +def json_safe_realtime_payload(payload: Mapping[str, object]) -> dict[str, object]: + clean: dict[str, object] = {} + for key, value in payload.items(): + if key == "extra_body": + continue + if isinstance(value, str | int | float | bool) or value is None: + clean[key] = value + elif isinstance(value, dict): + clean[key] = json_safe_realtime_payload(value) + elif isinstance(value, list): + clean[key] = [ + (json_safe_realtime_payload(item) if isinstance(item, dict) else item) + for item in value + if isinstance(item, str | int | float | bool | dict) or item is None + ] + return clean + + +def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: + transcription = session_payload.get("input_audio_transcription") + if isinstance(transcription, dict): + return transcription + audio_config = session_payload.get("audio") + if not isinstance(audio_config, dict): + return None + audio_input = audio_config.get("input") + if not isinstance(audio_input, dict): + return None + transcription = audio_input.get("transcription") + return transcription if isinstance(transcription, dict) else None + + +def realtime_max_output_tokens(value: object) -> int | None: + """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" + if value is None: + return None + if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: + return None + if isinstance(value, int) and value > 0: + return int(value) + return None From f5e4f5fa429ae3b5a8c50e83785900faea1ecc7b Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:35:46 +0800 Subject: [PATCH 28/99] [Bugfix][MiniCPM-o] Size the duplex HD slice reservation from the frame (#7654) Signed-off-by: Tianyao Wu Co-authored-by: Claude --- .../duplex/test_scheduler_token_budget.py | 245 ++++++++++++++++++ .../models/minicpmo_4_5/duplex/plugin.py | 111 ++++++-- 2 files changed, 339 insertions(+), 17 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py new file mode 100644 index 000000000..03072cb33 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""What a duplex append reserves in the scheduler, and why that number. + +The reservation is not advisory: ``build_duplex_data_plane_prompt`` turns it +into ``[scheduler_token_id] * budget``, and a unit whose embeddings outnumber +its slots has its tail dropped with a warning rather than failing +(``MiniCPMO45OmniModel.get_input_embeddings``). Over-reserving only wastes +slots, so every uncertainty here has to resolve upwards. + +The camera side depends on HD slicing, which depends on the frame size +*relative to the checkpoint's normalization tile*. The tile is configuration, +not a constant, so the cases below drive the arithmetic at more than one tile +size and check it against ``MiniCPMVImageProcessor.get_sliced_grid`` -- the +same grid search the checkpoint's own processor runs. +""" + +from __future__ import annotations + +import base64 +import io +from dataclasses import dataclass + +import pytest +from PIL import Image + +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import ( + PRIVATE_RUNTIME_CONFIG_KEYS, + _apply_default_scheduler_policy, + _duplex_vision_tile_pixels, + _duplex_vision_tokens, + _model_vision_tile_pixels, + duplex_scheduler_token_budget, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import MiniCPMVImageProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TOKENS_PER_BLOCK = 66 +#: ``_official_max_slice_nums`` gives the first frame of a stacked pair 2 and +#: every other frame 1. +BASE_FRAME_MAX_SLICES = 2 +#: What the released MiniCPM-o 4.5 checkpoint configures. +TILE_448 = 448 * 448 + +SIZES = [(224, 224), (320, 180), (448, 448), (448, 449), (640, 480), (960, 540), (1280, 720), (540, 960)] + + +def _frame(size: tuple[int, int], fmt: str = "JPEG") -> str: + buffer = io.BytesIO() + Image.new("RGB", size).save(buffer, format=fmt) + return base64.b64encode(buffer.getvalue()).decode() + + +def _blocks_the_model_makes(size: tuple[int, int], max_slice_nums: int, scale_resolution: int) -> int: + """Source tile plus the HD patches, straight from the model's grid search.""" + processor = MiniCPMVImageProcessor(scale_resolution=scale_resolution) + grid = processor.get_sliced_grid(image_size=size, max_slice_nums=max_slice_nums) + return 1 + grid[0] * grid[1] if grid else 1 + + +@pytest.mark.parametrize("scale_resolution", [336, 448, 560]) +@pytest.mark.parametrize("size", SIZES, ids=[f"{w}x{h}" for w, h in SIZES]) +@pytest.mark.parametrize("fmt", ["JPEG", "PNG"]) +def test_a_stacked_pair_reserves_what_the_model_slices_it_into( + size: tuple[int, int], fmt: str, scale_resolution: int +) -> None: + """The tile comes from the checkpoint, so the arithmetic has to follow it.""" + frame = _frame(size, fmt) + expected_blocks = _blocks_the_model_makes(size, BASE_FRAME_MAX_SLICES, scale_resolution) + _blocks_the_model_makes( + size, 1, scale_resolution + ) + + reserved = _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=scale_resolution**2) + + assert reserved == expected_blocks * TOKENS_PER_BLOCK + + +def test_a_small_frame_no_longer_reserves_three_hd_blocks() -> None: + """A 448x448 frame is one tile, so a stacked pair is two blocks, not four. + + The size-independent count reserved 264 tokens for this pair. 132 of them + were for patches the model never produces. + """ + frame = _frame((448, 448)) + + assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK + + +def test_one_tile_is_the_whole_slicing_decision() -> None: + """``ceil(w * h / tile)`` capped at 2: a single pixel over the tile slices.""" + + def pair(size: tuple[int, int]) -> int: + frame = _frame(size) + return _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448) + + assert pair((448, 448)) == 2 * TOKENS_PER_BLOCK + assert pair((448, 449)) == 4 * TOKENS_PER_BLOCK + + +def test_only_the_first_frame_of_a_pair_is_hd_sliced() -> None: + """``_official_max_slice_nums`` is ``[2, 1]``: the composite is never sliced. + + Sizing the wrong half of the pair would pass every same-size case, so the + two orders have to disagree. + """ + small, large = _frame((448, 448)), _frame((1280, 720)) + + assert _duplex_vision_tokens({"video_frames": [small, large]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [large, small]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK + + +def test_a_lone_frame_is_never_hd_sliced() -> None: + """``_official_max_slice_nums(1)`` is ``[1]``: one frame is one tile at any size.""" + assert _duplex_vision_tokens({"video_frames": [_frame((1280, 720))]}, tile_pixels=TILE_448) == TOKENS_PER_BLOCK + + +def test_extra_frames_past_the_pair_each_cost_one_block() -> None: + """The wire caps a frame list at two; the in-process API does not.""" + small, large = _frame((448, 448)), _frame((1280, 720)) + + assert _duplex_vision_tokens({"video_frames": [small] * 3}, tile_pixels=TILE_448) == 3 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [large] * 3}, tile_pixels=TILE_448) == 5 * TOKENS_PER_BLOCK + + +def test_an_unknown_tile_size_keeps_the_sliced_reservation() -> None: + """No tile, no shrinking. A checkpoint whose config cannot be read keeps the old number.""" + frame = _frame((224, 224)) + + assert _duplex_vision_tokens({"video_frames": [frame, frame]}) == 4 * TOKENS_PER_BLOCK + assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=None) == 4 * TOKENS_PER_BLOCK + + +@pytest.mark.parametrize( + "first_frame", + [ + pytest.param("data:image/jpeg;base64," + _frame((224, 224)), id="data url"), + pytest.param("!!! not base64 !!!", id="corrupt base64"), + pytest.param(base64.b64encode(b"\x00" * 64).decode(), id="base64 of something else"), + pytest.param(base64.b64encode(base64.b64decode(_frame((224, 224)))[:8]).decode(), id="truncated jpeg"), + ], +) +def test_a_frame_whose_header_will_not_parse_keeps_the_sliced_reservation(first_frame: str) -> None: + """The Realtime wire screens most of these; the in-process API does not. + + ``validate_realtime_video_frames`` rejects bad base64 and anything whose + magic bytes are not JPEG or PNG, so on a websocket session only a truncated + but well-headed frame gets this far. ``DuplexOmni.append_audio`` submits + without that screen, so the reservation still has to hold on its own. + """ + good = _frame((224, 224)) + + assert _duplex_vision_tokens({"video_frames": [first_frame, good]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK + + +def test_no_camera_track_reserves_nothing() -> None: + for payload in ({}, {"video_frames": []}, {"video_frames": "not a list"}, {"video_frames": [None, ""]}, None): + assert _duplex_vision_tokens(payload, tile_pixels=TILE_448) == 0 + + +def test_the_audio_budget_is_untouched_by_the_camera_track() -> None: + """One second of pcm_f32le is 12 slots; the frames are added on top.""" + one_second = base64.b64encode(b"\x00" * 4 * 16000).decode() + audio_only = {"audio": one_second, "format": "pcm_f32le"} + with_frames = {**audio_only, "video_frames": [_frame((448, 448)), _frame((448, 448))]} + + assert duplex_scheduler_token_budget(audio_only, tile_pixels=TILE_448) == 12 + assert duplex_scheduler_token_budget(with_frames, tile_pixels=TILE_448) == 12 + 2 * TOKENS_PER_BLOCK + + +# ---- where the tile comes from ---- + + +@dataclass(frozen=True) +class _SliceConfig: + scale_resolution: int + + +@dataclass(frozen=True) +class _HFConfig: + slice_config: object = None + image_size: object = None + + +@dataclass(frozen=True) +class _ModelConfig: + hf_config: object = None + + +@pytest.mark.parametrize( + ("hf_config", "expected"), + [ + (_HFConfig(slice_config={"max_slice_nums": 1, "scale_resolution": 448}, image_size=448), TILE_448), + (_HFConfig(slice_config=_SliceConfig(scale_resolution=336), image_size=448), 336 * 336), + (_HFConfig(image_size=560), 560 * 560), + (_HFConfig(), None), + (_HFConfig(slice_config={"scale_resolution": 0}, image_size=0), None), + ], + ids=["released-checkpoint", "attribute", "image-size-fallback", "nothing-to-read", "nonsense"], +) +def test_the_tile_is_read_from_the_checkpoint(hf_config: object, expected: int | None) -> None: + """First row is the released MiniCPM-o 4.5 ``config.json``, where ``slice_config`` is a dict.""" + assert _model_vision_tile_pixels(_ModelConfig(hf_config=hf_config)) == expected + + +def test_no_model_config_means_no_tile() -> None: + assert _model_vision_tile_pixels(None) is None + assert _model_vision_tile_pixels(_ModelConfig()) is None + + +def test_the_tile_reaches_the_budget_through_the_runtime_config() -> None: + """``_apply_default_scheduler_policy`` writes it; ``build_duplex_data_plane_prompt`` reads it back.""" + runtime_config: dict[str, object] = {} + + _apply_default_scheduler_policy( + runtime_config, + config=DuplexSessionConfig(), + tokenizer=None, + model_config=_ModelConfig(hf_config=_HFConfig(slice_config={"scale_resolution": 448})), + ) + + assert runtime_config["duplex_vision_tile_pixels"] == TILE_448 + assert _duplex_vision_tile_pixels(runtime_config) == TILE_448 + + +def test_a_checkpoint_that_cannot_be_read_leaves_the_key_out() -> None: + """No key, no tile, and no tile is the sliced reservation.""" + runtime_config: dict[str, object] = {} + + _apply_default_scheduler_policy(runtime_config, config=DuplexSessionConfig(), tokenizer=None) + + assert "duplex_vision_tile_pixels" not in runtime_config + assert _duplex_vision_tile_pixels(runtime_config) is None + + +def test_a_junk_tile_in_the_runtime_config_is_ignored() -> None: + for value in (0, -1, "448", 448.0, None): + assert _duplex_vision_tile_pixels({"duplex_vision_tile_pixels": value}) is None + + +def test_the_tile_is_server_owned() -> None: + """A client that could set the tile could shrink the reservation under the worker.""" + assert "duplex_vision_tile_pixels" in PRIVATE_RUNTIME_CONFIG_KEYS diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py index 673640461..e57377766 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -63,6 +63,15 @@ _DUPLEX_VISION_TOKENS_PER_FRAME = 66 # Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD # sliced (1 source + 2 patches on 960x540) and the composite is not. _DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 +# How many blocks that slicing costs depends on the frame, and at +# ``max_slice_nums=2`` the processor's grid search has only two outcomes. +# ``MiniCPMVImageProcessor.get_sliced_grid`` takes +# ``multiple = min(ceil(w * h / scale_resolution**2), max_slice_nums)`` and +# returns no grid at all for ``multiple <= 1``; at 2 the only candidate split +# is 2, so a frame is either unsliced or a 2-cell grid. One normalization tile +# is therefore the whole decision -- but ``scale_resolution`` is the +# checkpoint's, not a constant, so it is read from the model rather than +# assumed. PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( { @@ -70,6 +79,7 @@ PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( "duplex_stage_max_tokens", "duplex_stage0_max_tokens", "duplex_scheduler_token_id", + "duplex_vision_tile_pixels", "duplex_first_append_context_tokens", "ref_audio_data", "ref_audio_format", @@ -86,29 +96,68 @@ class MiniCPMO45ClientRuntimeConfigError(DuplexRuntimeConfigError): # ---- engine policy helpers: scheduler token budget ---- -def _duplex_frame_count(payload: object) -> int: +def _duplex_frames(payload: object) -> list[str]: if not isinstance(payload, dict): - return 0 + return [] frames = payload.get("video_frames") if not isinstance(frames, list): - return 0 - return sum(1 for frame in frames if isinstance(frame, str) and frame) + return [] + return [frame for frame in frames if isinstance(frame, str) and frame] -def _duplex_vision_tokens(payload: object) -> int: +def _duplex_base_frame_blocks(frame: str, tile_pixels: int | None) -> int: + """Blocks the HD-sliced frame of a stacked pair costs, read from the frame. + + Wire format is the one Stage0 decodes: bare base64 JPEG/PNG, and only the + header is parsed -- the pixels are the worker's job. + + Every uncertainty resolves to the sliced count, because the two directions + are not symmetric. Over-reserving wastes scheduler slots. Under-reserving + hands the worker fewer prompt slots than it has embeddings, and + ``MiniCPMO45OmniModel`` then drops the tail of the unit with a warning + rather than failing, which is silent audio loss. + """ + if tile_pixels is None: + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + + from io import BytesIO + + from PIL import Image + + try: + raw = b64decode(frame, validate=True) + with Image.open(BytesIO(raw)) as image: + width, height = image.size + except (BinasciiError, ValueError, OSError, Image.DecompressionBombError): + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + if width * height <= tile_pixels: + return 1 + return _DUPLEX_HD_SLICES_PER_BASE_FRAME + + +def _duplex_vision_tile_pixels(runtime_config: object) -> int | None: + """Area of the tile this model normalizes a frame to, or ``None`` if unknown.""" + if not isinstance(runtime_config, dict): + return None + value = runtime_config.get("duplex_vision_tile_pixels") + return value if isinstance(value, int) and value > 0 else None + + +def _duplex_vision_tokens(payload: object, *, tile_pixels: int | None = None) -> int: """Scheduler slots for this append's camera track. Audio is never stacked: a unit still carries one second of soundtrack. ``stack_frames`` only adds a second *image*. Official HD on that pair is - ``[2, 1]``, so the base frame reserves three 66-token blocks and every - extra frame reserves one. + ``[2, 1]``, so only the first frame is sliced and every other frame + reserves one 66-token block. """ - count = _duplex_frame_count(payload) - if count <= 0: + frames = _duplex_frames(payload) + if not frames: return 0 - if count >= 2: - return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME - return count * _DUPLEX_VISION_TOKENS_PER_FRAME + if len(frames) == 1: + return _DUPLEX_VISION_TOKENS_PER_FRAME + blocks = _duplex_base_frame_blocks(frames[0], tile_pixels) + (len(frames) - 1) + return blocks * _DUPLEX_VISION_TOKENS_PER_FRAME def _duplex_pcm_sample_count(payload: object) -> int | None: @@ -136,8 +185,8 @@ def duplex_first_append_unit_count(payload: object) -> int | None: return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) -def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: - vision_tokens = _duplex_vision_tokens(payload) +def duplex_scheduler_token_budget(payload: object, *, default: int = 64, tile_pixels: int | None = None) -> int: + vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels) sample_count = _duplex_pcm_sample_count(payload) if sample_count is None: return max(1, int(default)) + vision_tokens @@ -185,13 +234,15 @@ def build_duplex_data_plane_prompt( payload: object, final: bool, ) -> dict[str, object]: - token_budget = duplex_scheduler_token_budget(payload) + tile_pixels = _duplex_vision_tile_pixels(runtime_config) + token_budget = duplex_scheduler_token_budget(payload, tile_pixels=tile_pixels) if seq <= 1: context_reserve = duplex_first_append_context_reserve(runtime_config) token_budget += context_reserve first_units = duplex_first_append_unit_count(payload) if first_units is not None: - token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) + vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels) + token_budget = context_reserve + first_units * 12 - 1 + vision_tokens if seq > 1 and duplex_payload_is_exact_chunks(payload): token_budget += 1 if final and duplex_payload_is_exact_chunks(payload): @@ -466,11 +517,34 @@ def _apply_first_append_context_tokens( runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) +def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None: + """Area of one normalization tile, from the checkpoint that will do the slicing. + + ``MiniCPMVImageProcessor`` is built with ``scale_resolution=config.image_size`` + and Stage0 loads the checkpoint's own processor, so this is per-checkpoint + configuration. ``None`` when it cannot be read, which keeps the reservation + at the sliced count. + """ + hf_config = getattr(model_config, "hf_config", None) + if hf_config is None: + return None + slice_config = getattr(hf_config, "slice_config", None) + side = getattr(slice_config, "scale_resolution", None) + if not isinstance(side, int): + side = slice_config.get("scale_resolution") if isinstance(slice_config, dict) else None + if not isinstance(side, int): + side = getattr(hf_config, "image_size", None) + if not isinstance(side, int) or side <= 0: + return None + return side * side + + def _apply_default_scheduler_policy( runtime_config: dict[str, object], *, config: DuplexSessionConfig, tokenizer: PreTrainedTokenizerBase | None, + model_config: ModelConfig | None = None, ) -> None: stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} @@ -491,6 +565,9 @@ def _apply_default_scheduler_policy( scheduler_token_id = _scheduler_token_id(tokenizer) if scheduler_token_id is not None: runtime_config["duplex_scheduler_token_id"] = scheduler_token_id + tile_pixels = _model_vision_tile_pixels(model_config) + if tile_pixels is not None: + runtime_config["duplex_vision_tile_pixels"] = tile_pixels class MiniCPMO45DuplexPlugin(DuplexModelPlugin): @@ -653,7 +730,7 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): if isinstance(initial_user_text, str) and initial_user_text: runtime_config["initial_user_text"] = initial_user_text tokenizer = await self._tokenizer_for(model_config) - _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer) + _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer, model_config=model_config) ref_audio = config.ref_audio extra_ref_audio = extra_body.get("ref_audio") From bf74544a35bb45d6092fb2b568e01fc51859b09b Mon Sep 17 00:00:00 2001 From: Clodagh Walsh <42538837+clodaghwalsh17@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:56:58 +0100 Subject: [PATCH 29/99] [Voxtral] Follow up to raise BadRequest for ref_audio (#7649) Signed-off-by: Clodagh Walsh --- vllm_omni/entrypoints/openai/tts_adapters/voxtral.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py index 5ea82babd..4cdf990f2 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py @@ -99,9 +99,10 @@ class VoxtralTTSAdapter(ARTTSAdapter): ) -> PreparedRequest: # Cache encoder_loaded property to avoid RPC call on every request with ref_audio if self._encoder_loaded is None: - loaded = await self.ctx.server.engine_client.collective_rpc("encoder_loaded") - # Flatten results before checking if encoder loaded property is set - self._encoder_loaded = any(sum(loaded, [])) + result = await self.ctx.server.engine_client.collective_rpc("encoder_loaded") + self._encoder_loaded = any( + bool(x) for r in result if not isinstance(r, dict) for x in (r if isinstance(r, list) else [r]) + ) prompt = await self._build_prompt_async(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts") From 873e9bff7c545c5cda79fdb93a867f09e443b61d Mon Sep 17 00:00:00 2001 From: ooooooye <53851298+brandneway@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:05:15 +0800 Subject: [PATCH 30/99] [Doc][NPU] Refresh MiniMax-H3 Ascend recipes for 0.28-era configurations (#7720) Signed-off-by: brandneway --- recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md | 337 ++++++++-------- recipes/MiniMaxAI/MiniMax-H3-NPU.md | 456 ++++++++++++++-------- 2 files changed, 477 insertions(+), 316 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md index 36eb78a5d..039b40272 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md @@ -1,195 +1,224 @@ -# MiniMax-H3 on a single NPU 950PR +# MiniMax H3 on Ascend 950PR / 950DT -This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR -NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at -1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see -[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md). +> Joint video and audio generation — Ascend NPU deployment guide for the +> four-card Ascend 950PR / 950DT route -## Capacity requirements +## Summary -| Resource | Requirement | -| --- | ---: | -| NPU | 1x NPU 950PR | -| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) | -| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) | -| Checkpoint storage | 135 GiB per partition, local disk strongly preferred | -| Container shared memory | 8 GiB, or a writable local filesystem for spill | - -The observed high-water mark leaves about 12.3 GiB of headroom, so this card is -sized for exactly one resident task partition at 1024x576. Do not expect a -larger output shape or concurrency greater than one to fit. - -Because a single 950PR holds the whole partition in HBM, this route does not -need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB -GPU recipes does not apply. The validated run completed inside a container -limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is -actively harmful: it stages weights in non-reclaimable host memory and the -container's OOM killer terminates the server with exit code -9. +- Vendor: MiniMaxAI +- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) +- Tasks: T2VA, FL2VA, and Ref2VA +- Mode: OpenAI-compatible `/v1/videos` HTTP serving +- Hardware: Ascend 950PR / Ascend 950DT, 4x NPU (128 GB HBM per device) +- Maintainer: Community +This recipe covers the four-card 950PR / 950DT configuration at 1344x768. +For the eight-card Atlas 800I A2 / A3 route, see +[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md); checkpoint layout, container +preparation, MindIE-SD installation, and the ffmpeg/decord dependencies are +described there and apply unchanged. ## Environment -- Host architecture: x86_64 -- driver: 25.7.rc1.6 (hal 7.35.23) -- firmware: 9.0.0.105.229 -- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`) -- npu-smi: 25.7.rc1.6 -- Python: 3.12.13 -- PyTorch: 2.10.0+cpu -- torch_npu: 2.10.0.post2 -- vLLM: 0.26.0 -- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`) - -Install vLLM-Omni from a checkout with MiniMax-H3 support: +Use the A5 variant of the official vLLM-Omni NPU images (see +[NPU installation](../../docs/getting_started/installation/npu.md)): ```bash -uv venv -source .venv/bin/activate -uv pip install -e . +export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a5 ``` +Check the [tag list](https://quay.io/repository/ascend/vllm-omni?tab=tags) +for newer releases. Inside the container, install MindIE-SD and the optional +Ref2VA media dependencies as described in +[MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment). + ## Start a server +Task selection follows the eight-card recipe: pass `--task-type fl2va` or +`--task-type ref2va` to pick the task partition. The configurations below use +4-card USP, 4-card text-encoder TP, and VAE `tile` parallelism. + +The two recommended configurations differ in how they fit memory: the +lossless configuration enables distributed layerwise offload (DLO) to avoid +activation OOM at longer durations, while the lossy configuration fits the +maximum supported MiniMax-H3 workload shape without OOM, so DLO is omitted +for performance. + +### Recommended lossless configuration (Ascend 950PR / 950DT) + ```bash -export MODEL=/path/to/MiniMax-H3/FL2VA -export PORT=8000 +export PORT=9098 +export MODEL=/path/to/MiniMax-H3 export VLLM_WORKER_MULTIPROC_METHOD=spawn -export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 +export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=4000 +export PYTHONDONTWRITEBYTECODE=1 +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export HCCL_NPU_SOCKET_PORT_RANGE="auto" vllm serve "${MODEL}" \ --omni \ --host 0.0.0.0 \ --port "${PORT}" \ --trust-remote-code \ - --num-gpus 1 \ - --tensor-parallel-size 1 \ - --usp 1 \ + --task-type fl2va \ + --num-gpus 4 \ + --usp 4 \ --ring 1 \ - --text-encoder-tp-size 1 \ - --vae-patch-parallel-size 1 \ + --text-encoder-tp-size 4 \ --vae-parallel-mode tile \ --vae-use-tiling \ - --init-timeout 14400 \ - --stage-init-timeout 14400 \ - --quantization int8 + --vae-patch-parallel-size 4 \ + --enable-diffusion-pipeline-profiler \ + --enable-distributed-layerwise-offload \ + --diffusion-attention-backend FLASH_ATTN ``` -Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no -NPU implementation and fails at the first denoise step. Leave the backend -unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is -installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)). +H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. The first +request includes regional compilation; warm the server once before measuring +steady-state latency. -H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. +### Recommended lossy configuration (Ascend 950PR / 950DT) -## Validated evidence - -Measured on one NPU 950PR with the server command above plus -`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from -the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3 -requested 60 denoise steps and executed 59 denoise updates, so per-step latency -is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps, -124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples. - -Five requests were issued back to back with no idle time between them. The card -is measurably faster on the first request than in steady state, so both regimes -are reported. - -| Stage | First request, cold card | Thermal steady state | -| --- | ---: | ---: | -| End-to-end request | 462.13 s | 506.38 s | -| Text encode | 0.55 s | 0.05 s | -| Denoise (59 updates) | 429.29 s | 471.74 s | -| Per denoise update | 7,276 ms | 7,996 ms | -| VAE decode | 8.43 s | 8.14 s | -| Worker-to-server handoff | 24 s | 27 s | -| MP4 muxing | 0.57 s | 0.50 s | -| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms | -| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB | -| Average NPU power | 566.9 W | 556.9 W | -| Peak NPU temperature | 104 C | 104 C | - -Sample counts differ per row. The cold column is a single request. In the steady -column, text encode, denoise and VAE decode are means over the four subsequent -requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while -end-to-end, handoff and muxing are means over the two of those four that -survived the handoff timeout described above. - -**Report the steady-state column.** The cold card completes denoise 9.9% faster -while drawing only 1.8% more average power, and both regimes reach the same -104 C ceiling, so the first request is buying a short window of higher clocks -before the die saturates. Any benchmark that issues a single request against an -idle card will overstate throughput by about 10%. - -**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides -the entire stage wall time by the 60 *requested* steps, so it absorbs text -encoding, VAE decode and the 837 MiB handoff. All three successful requests -satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The -per-update figure in the table is `diffuse / 59` from the pipeline profiler. - -Denoise accounts for 93% of the steady-state request and the handoff accounts -for 5.3%, so keeping the transfer in shared memory should bring end-to-end down -to roughly 480 s. - -Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a -reserved high-water mark rather than live allocation: it reads 111,272 MiB after -weight loading and before the first request, and does not fall back after -requests complete. Re-measure for longer outputs, a different output shape, or -concurrency greater than one. - -### Startup - -Weight loading took **31 min 50 s** on the validated host (first log line -08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a -shared GlusterFS network mount. This is why the server command sets -`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and -1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143. -Reduce both timeouts when the partition is staged on local disk. - -## T2VA request example +Adds EQBSA sparse attention (block-sparse with Q/K INT8 + V FP8 mixed +precision) and MXFP8 online quantization, and drops DLO: ```bash -export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync" +export PORT=9098 +export MODEL=/path/to/MiniMax-H3 +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=4000 +export PYTHONDONTWRITEBYTECODE=1 +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export HCCL_NPU_SOCKET_PORT_RANGE="auto" -curl -sS --max-time 1800 -X POST "${API_URL}" \ - -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \ - -F 'width=1024' \ - -F 'height=576' \ - -F 'aspect_ratio=16:9' \ - -F 'fps=24' \ - -F 'num_inference_steps=60' \ - -F 'flow_shift=12' \ - -F 'seed=1101' \ - -F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \ - -o t2va.mp4 +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --task-type fl2va \ + --num-gpus 4 \ + --usp 4 \ + --ring 1 \ + --text-encoder-tp-size 4 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --vae-patch-parallel-size 4 \ + --enable-diffusion-pipeline-profiler \ + --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN","block_sparse":{"sparsity":0.8,"precision":"mix","start_step":8,"end_step":12}}}' \ + --diffusion-quantization-config '{"transformer":{"method":"mxfp8"}}' ``` +The recommended lossy configuration uses a **start=8 / end=12** dense-fallback +window; see below for more aggressive settings. + +### Optional optimizations + +The following sections list only the **increment or replacement** relative to +the recommended lossless configuration; environment variables and the +remaining flags stay unchanged. + +#### Tuning the EQBSA dense-fallback window + +`start_step` and `end_step` are the numbers of leading/trailing steps that +fall back to dense attention — not step indices. `precision` defaults to +`bf16`; EQBSA requires setting it to `mix` explicitly. For 1344x768, 50-step +T2VA: + +| Configuration | `start_step` | `end_step` | Notes | +| --- | ---: | ---: | --- | +| Quality first (recommended) | 8 | 13 | General and complex-motion scenes | +| Balanced | 0 | 13 | Balanced | +| Speed first | 0 | 0 | Simple, low-motion scenes | + +When frames look discontinuous or details unstable, increase `end_step` +first, then `start_step`; do not compensate by lowering `sparsity` alone. +These pairings were validated at 1344x768 / 50 steps only; re-evaluate for +other resolutions or step counts. + +#### Cache-DiT + +Append to enable DiT block caching with TaylorSeer extrapolation — see the +[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md): + +```bash + --cache-backend cache_dit \ + --enable-cache-dit-summary \ + --cache-config '{"Fn_compute_blocks":2,"Bn_compute_blocks":1,"max_warmup_steps":4,"residual_diff_threshold":0.4,"max_continuous_cached_steps":4,"enable_taylorseer":true,"taylorseer_order":2}' +``` + +#### MXFP4 online quantization + +Append to switch the DiT to MXFP4 online quantization: + +```bash + --diffusion-quantization-config '{"transformer":{"method":"mxfp4"}}' +``` + +and remove `--enable-distributed-layerwise-offload`. + +## Request examples + +Identical to the eight-card recipe; see +[MiniMax-H3-NPU.md § HTTP API examples](MiniMax-H3-NPU.md#http-api-examples) +and the GPU recipe's full parameter table. + +## Benchmarks (Ascend 950PR) + +Measured with `--enable-diffusion-pipeline-profiler` on vLLM-Omni 0.28.0: +lossless rows use the recommended lossless configuration (with DLO), lossy +rows the recommended lossy configuration (EQBSA + MXFP8, without DLO). Memory +figures are per device. + +| Task | Resolution | Frames | Duration (s) | Config | E2E (s) | DiT total (s) | DiT per-step (s) | DiT steps | Text encode (s) | Ref video encode (s) | Ref audio encode (s) | VAE decode (s) | Ref preprocess (s) | Post-process (s) | CPU MP4 (s) | Resident weights (GB) | Peak memory (GB) | +| ---- | ---- | ---- | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| t2va | 1344x768 | 124 | 5 | lossless | 203.95 | 198.10 | 4.04 | 49 | 0.57 | NA | NA | 5.05 | NA | 0.001 | 0.23 | 3.97 | 22.2 | +| t2va | 1344x768 | 362 | 15 | lossless | 1489.17 | 1474.35 | 30.08 | 49 | 0.69 | NA | NA | 13.45 | NA | 0.001 | 0.82 | 3.97 | 32.14 | +| t2va | 1344x768 | 124 | 5 | lossy | 113.86 | 109.81 | 2.24 | 49 | 0.03 | NA | NA | 3.82 | NA | 0.001 | 0.21 | 56.99 | 64.38 | +| t2va | 1344x768 | 362 | 15 | lossy | 762.17 | 749.40 | 15.29 | 49 | 0.03 | NA | NA | 12.19 | NA | 0.001 | 0.61 | 56.99 | 74.26 | +| ref2va | 1344x768 | 124 | 5 | lossless | 745.43 | 730.22 | 14.9 | 49 | 1.16 | 8.15 | 0.13 | 5.12 | 0.41 | 0.001 | 0.28 | 3.97 | 23.82 | +| ref2va | 1344x768 | 362 | 15 | lossless | 5887.24 | 5847.57 | 119.33 | 49 | 2.54 | 22.07 | 0.23 | 13.34 | 0.81 | 0.001 | 0.70 | 3.97 | 32.53 | +| ref2va | 1344x768 | 124 | 5 | lossy | 400.40 | 387.82 | 7.91 | 49 | 0.95 | 7.05 | 0.09 | 3.91 | 0.41 | 0.001 | 0.21 | 56.6 | 66.03 | +| ref2va | 1344x768 | 362 | 15 | lossy | 3037.84 | 3001.31 | 61.24 | 49 | 2.28 | 20.64 | 0.15 | 12.1 | 0.81 | 0.001 | 0.70 | 56.6 | 74.76 | + +The resident-weight figures reflect the DLO strategy difference: with DLO +(lossless rows) DiT weights live in host memory and only the active layer +stays on the device, while without DLO (lossy rows) the quantized weights are +resident in HBM. Peak memory stays well inside the 128 GB per-device capacity +in both cases. These numbers describe the validated shapes rather than a +general throughput guarantee. + ## Known limitations -- INT8 online quantization is validated for T2VA on this card. Use BF16 for - FL2VA and Ref2VA, or re-measure before relying on it. -- Single-card serving loads one task partition at a time. For Ref2VA, stop the - server and restart it against the `Ref2VA` directory. -- Sustained load drives the die to 104 C and costs about 10% of denoise - throughput relative to a cold card. Warm the server with at least one full - request before measuring, and treat steady state as the reportable number. -- Requests can fail during result handoff on hosts where the IPC spill path is - slow. See [§ Result handoff and shared - memory](#result-handoff-and-shared-memory). -- Loading the checkpoint from a network mount takes over 30 minutes and forces - very large init timeouts. Stage the partition on local disk when possible. -- The configuration measured here is 1024x576 at 60 steps, which differs from - the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise - cost scales superlinearly with token count, so the numbers are not directly - comparable across recipes. -- The image ships a `triton` package whose backend is not built - (`No module named 'triton._C.libtriton.'`). vLLM logs this as an error - at startup and disables Triton; the diffusion path does not need it and the - run completes normally. +- Task serving is partitioned by `--task-type`: T2VA/FL2VA and Ref2VA load + different DiTs, so switching between them requires a restart with the other + partition. +- H3 currently executes one generation request per diffusion batch. +- The first regional-compile request is a warmup and should not be included + in steady-state performance measurements. +- VAE patch parallelism requires size 1 or the full DiT group size and + supports the H3 native `tile` mode only. +- The lossy configuration (EQBSA + MXFP8) is validated for T2VA and Ref2VA; + use the lossless configuration for FL2VA or re-validate first. +- The lossless configuration trades throughput for activation headroom via + DLO; if your workload never approaches the maximum duration, dropping + `--enable-distributed-layerwise-offload` may improve latency at the risk of + OOM on longer durations. ## Additional resources - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide -- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide -- [Int8 quantization](../../docs/user_guide/quantization/int8.md) +- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A2 / A3 + guide (checkpoint, container, and dependency details) +- [NPU installation](../../docs/getting_started/installation/npu.md) +- [RainFusion attention](../../docs/user_guide/diffusion/attention_backends/rainfusion.md) + — block-sparse knobs and tuning +- [Online quantization](../../docs/user_guide/quantization/online.md) + — INT8 / MXFP8 / MXFP4 +- [Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) - [Supported models](../../docs/models/supported_models.md) - [Video API](../../docs/serving/videos_api.md) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index c6ee0904b..971be278d 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -1,7 +1,8 @@ # MiniMax H3 on Ascend NPU > Joint video and audio generation with text, first/last keyframes, and mixed -> image/video/audio references — Ascend NPU deployment guide +> image/video/audio references — Ascend NPU deployment guide for Atlas 800I A2 +> and Atlas 800I A3 ## Summary @@ -9,7 +10,7 @@ - Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) - Tasks: T2VA, FL2VA, and Ref2VA - Mode: OpenAI-compatible `/v1/videos` HTTP serving -- Hardware: Atlas 800I A3 +- Hardware: Atlas 800I A2 / Atlas 800I A3, 8x NPU - Maintainer: Community This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU @@ -22,103 +23,288 @@ environments. Differences from the GPU path: automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio inputs. +For the four-card Ascend 950PR / 950DT route, see +[MiniMax-H3-NPU-950PR.md](MiniMax-H3-NPU-950PR.md). + ## Prerequisites ### Checkpoint -Same as the GPU recipe — Hugging Face access approval is required. Authenticate -once; `vllm serve` downloads the required nested components automatically: +Download the weights once; `vllm serve` accepts the repository ID or a local +path and resolves the nested components automatically: ```bash -hf auth login -export MODEL=MiniMaxAI/MiniMax-H3 +export MODEL=/path/to/MiniMax-H3 +hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL}" +``` + +The checkpoint directory layout is: + +```text +MiniMax-H3/ +├── FL2VA/ # shared by T2VA/FL2VA tasks +├── Ref2VA/ # Ref2VA task weights +└── model_index.json ``` ### Environment -- Ascend driver & firmware: 25.5.1 -- CANN toolkit: 9.0.1 -- Python: 3.12 -- PyTorch: 2.10.0+cpu -- torch_npu: 2.10.0.post2 -- Install vLLM-Omni from a checkout with MiniMax H3 support: +The recommended deployment uses the official vLLM-Omni NPU images, which ship +the matching vLLM / vLLM-Ascend / vLLM-Omni stack (see +[NPU installation](../../docs/getting_started/installation/npu.md)). Pick the +image matching the hardware — `quay.io/ascend/vllm-omni:v0.29.0` for +Atlas 800I A2, `quay.io/ascend/vllm-omni:v0.29.0-a3` for Atlas 800I A3 — and +check the +[tag list](https://quay.io/repository/ascend/vllm-omni?tab=tags) for newer +releases: ```bash -uv venv -source .venv/bin/activate -uv pip install -e . +export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a3 +export CONTAINER_NAME=h3 + +docker run -it -u root --name ${CONTAINER_NAME} \ + --privileged=true \ + --shm-size=2000g \ + --net=host \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ + --device /dev/davinci4 \ + --device /dev/davinci5 \ + --device /dev/davinci6 \ + --device /dev/davinci7 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /home:/home \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v /root/.cache:/root/.cache \ + ${IMAGE} \ + /bin/bash ``` -Install the **mindie-sd** operator library to enable Ascend-optimized fused -operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse -attention kernel: +Inside the container: + +- **MindIE-SD** provides the Ascend-optimized fused operators + (`adalayernorm`, etc.), the LaserAttention kernel, and the RainFusion + block-sparse attention kernel. If the image does not ship it, build it from + source: ```bash -git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD - -# Comment out the tik_ops build step (not needed for this use case) -sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh - +git clone https://gitcode.com/Ascend/MindIE-SD.git +cd MindIE-SD python setup.py bdist_wheel -cd dist -pip install mindiesd-*.whl +pip install dist/mindiesd-*.whl ``` -- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for - reference-video preparation and MP4 output. -- Reference-video decoding uses `decord` when available and falls back to - PyAV otherwise. -- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded - at native sample rate through the soundfile / ffmpeg fallback. +- To track the latest vLLM-Omni, install it from a checkout: + +```bash +git clone https://github.com/vllm-project/vllm-omni.git +cd vllm-omni +VLLM_OMNI_TARGET_DEVICE=npu pip install -e . --no-build-isolation +``` + +- `ffmpeg` and `ffprobe` must be available on `PATH` (reference-video + preparation and MP4 output). Reference-video decoding uses `decord` when + available and falls back to PyAV otherwise; `decord` is required only for + the Ref2VA workflow and can be built with: + +```bash +apt update +apt install -y ffmpeg pkg-config libavcodec-dev libavformat-dev \ + libavutil-dev libswscale-dev libavfilter-dev libavdevice-dev +git clone --depth 1 --recursive https://github.com/dmlc/decord.git +cd decord +mkdir -p build && cd build +cmake .. -DUSE_CUDA=0 -DDECODE_FFMPEG=1 -DCMAKE_BUILD_TYPE=Release +make +cd ../python +pip3 install . --no-build-isolation +``` ## Start a server -Pass the repository ID directly. The pipeline loads the two nested DiTs while -sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`. +T2VA and FL2VA share one set of weights; the Ref2VA weights differ. Select +the task partition at startup with **`--task-type`** (`fl2va` or `ref2va`). +When unset, only the T2VA/FL2VA weights are loaded. All remaining flags are +task-agnostic. -### Multi-NPU: 768P combined-service configuration +### Recommended lossless configuration (Atlas 800I A2 / A3) -Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism -degree 8, native tiled VAE patch parallelism degree 8, and distributed -layerwise offload: +Lossless means fused operators plus the high-performance FLASH_ATTN backend +only. The configuration uses 8-card USP, 8-card text-encoder TP, distributed +layerwise offload (DLO), and VAE `tile` parallelism: ```bash -export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 export PORT=9098 -export MODEL=MiniMaxAI/MiniMax-H3 +export MODEL=/path/to/MiniMax-H3 export VLLM_WORKER_MULTIPROC_METHOD=spawn -export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 +export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=4000 export PYTHONDONTWRITEBYTECODE=1 +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export MINDIE_SD_FA_TYPE="ascend_laser_attention" +export HCCL_NPU_SOCKET_PORT_RANGE="auto" vllm serve "${MODEL}" \ --omni \ --host 0.0.0.0 \ --port "${PORT}" \ --trust-remote-code \ + --task-type fl2va \ --num-gpus 8 \ --usp 8 \ --ring 1 \ --text-encoder-tp-size 8 \ - --enable-distributed-layerwise-offload \ --vae-parallel-mode tile \ --vae-use-tiling \ --vae-patch-parallel-size 8 \ + --enable-distributed-layerwise-offload \ + --enable-diffusion-pipeline-profiler \ --diffusion-attention-backend FLASH_ATTN ``` -Do not add `--enforce-eager`. The first request includes regional -compilation; warm the server once before measuring steady-state latency. -H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. +Notes: -The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no -partition restart is required. Layerwise offload applies to both DiTs. +- Do not add `--enforce-eager`. The first request includes regional + compilation; warm the server once before measuring steady-state latency. +- H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. +- Keep `--ring 1` when switching to RainFusion attention (below): the + block-sparse kernel ranks key blocks over the whole sequence, so ring + parallelism would split away the keys it needs. Scale with `--usp` instead. +- DLO offloads DiT layers with parameters gathered across the parallel group + instead of replicated per rank, which is what fits the combined 768P + service into 64 GB HBM per device (Atlas 800I A3). Host (CPU) memory usage + is high with this option. -On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at -768P without offloading or sharding: use distributed layerwise offload (as -above) or HSDP — see -[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3). +### Recommended lossy configuration (Atlas 800I A2 / A3) -### CPU MP4 response encoding (Atlas A2) +Adds RainFusion block-sparse attention and INT8 online quantization on top of +the lossless configuration. Validated for T2VA and Ref2VA; see +[Benchmarks](#benchmarks-atlas-800i-a2--a3): + +```bash +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --task-type fl2va \ + --num-gpus 8 \ + --usp 8 \ + --ring 1 \ + --text-encoder-tp-size 8 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --vae-patch-parallel-size 8 \ + --enable-distributed-layerwise-offload \ + --enable-diffusion-pipeline-profiler \ + --diffusion-quantization-config '{"transformer":{"method":"int8"}}' \ + --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN", + "block_sparse": {"sparsity": 0.8, "start_step": 12}}}' +``` + +Online INT8 quantization can be combined with distributed layerwise offload +while AllGather is enabled. Keep AllGather enabled for SP>1; use +`--dlo-no-use-allgather` only when collective communication is unavailable. + +### Optional optimizations + +The following sections list only the **increment or replacement** relative to +the recommended lossless configuration; environment variables and the +remaining flags stay unchanged. + +#### Cache-DiT + +Append to enable DiT block caching with TaylorSeer extrapolation — see the +[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md): + +```bash + --cache-backend cache_dit \ + --enable-cache-dit-summary \ + --cache-config '{"Fn_compute_blocks":2,"Bn_compute_blocks":1,"max_warmup_steps":4,"residual_diff_threshold":0.4,"max_continuous_cached_steps":4,"enable_taylorseer":true,"taylorseer_order":2}' +``` + +#### EQBSA sparse attention + +EQBSA applies Q/K INT8 with V FP8 mixed quantization on top of block-sparse +attention. Replace `--diffusion-attention-backend FLASH_ATTN` with: + +```bash + --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN","block_sparse":{"sparsity":0.8,"start_step":8,"end_step":13,"precision":"mix"}}}' +``` + +`precision` defaults to `bf16`; enabling EQBSA requires setting it to `mix` +explicitly. `start_step` and `end_step` are the numbers of leading/trailing +steps that fall back to dense attention — not step indices. For 1344x768, +50-step T2VA: + +| Configuration | `start_step` | `end_step` | Notes | +| --- | ---: | ---: | --- | +| Quality first (recommended) | 8 | 13 | General and complex-motion scenes | +| Balanced | 0 | 13 | Balanced | +| Speed first | 0 | 0 | Simple, low-motion scenes | + +When frames look discontinuous or details unstable, increase `end_step` +first, then `start_step`; do not compensate by lowering `sparsity` alone. +These pairings were validated at 1344x768 / 50 steps only; re-evaluate for +other resolutions or step counts. + +## HTTP API examples + +The request contract is identical to the GPU recipe; see +[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples). Use the +validated 768P shapes (e.g. `width=1344 height=768`) on NPU. Quick start +against the synchronous endpoint: + +```bash +export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync" + +curl -sS -X POST "${API_URL}" \ + -F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \ + -F 'width=1344' \ + -F 'height=768' \ + -F 'aspect_ratio=16:9' \ + -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'flow_shift=12.0' \ + -F 'audio_flow_shift=3.0' \ + -F 'seed=1101' \ + -F 'extra_params={"task":"t2va","duration":5.0}' \ + -o t2va.mp4 +``` + +Switch tasks via `extra_params.task` (`t2va` / `fl2va` / `ref2va`); FL2VA and +Ref2VA additionally pass `input_references` (first-frame image, reference +video/audio). The full parameter table lives in the GPU recipe. + +## FlashGen 4-step online LoRA (T2VA) + +The FlashGen 4-step weights are a native-layout LoRA single file +(`key_format=minimax-h3-native`) loaded at runtime — no merge into the base +checkpoint. Download them from +[FlashGen/Minimax-H3-4step-lora-flashgen](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen), +then on top of the recommended configuration above: + +- append `--lora-backend peft --lora-path "${FLASHGEN_LORA}"`, +- use `--diffusion-attention-backend FLASH_ATTN` (dense; drop the RainFusion + attention config), +- issue T2VA requests with `num_inference_steps=4` and the `lora` request + field (`name` / `path` / `scale`) — see + [MiniMax-H3.md § FlashGen native LoRA](MiniMax-H3.md#flashgen-native-lora). + +T2VA only; FL2VA and Ref2VA are not supported, and neither +`--enable-cpu-offload` nor `--enable-layerwise-offload` may be combined with +it. + +## CPU MP4 response encoding (Atlas A2) Non-streaming MP4 responses use one public automatic encoder. It checks the runtime frame shape, common dtype, and RGB channel-plane contiguity for every @@ -127,121 +313,65 @@ inputs fall back to the legacy muxer before opening the PyAV container. No CLI flag, model declaration, or user configuration is required. Streaming fMP4 output remains on its existing incremental path. -Current performance and correctness validation is limited to one Atlas A2 host -with 8x Ascend 910B4-1 NPUs, the `FL2VA`/`t2va` partition, one request at a -time, 1344x768 at 24 fps, and 5, 8.7, and 15 second requests. This optimization -only changes CPU MP4 response encoding; it does not change DiT execution or -stage 0. +## Ascend-optimized components -### Optional optimizations +The optimizations stacked by the recommended configurations above: -Two independent optimizations may be enabled on top of the configuration -above. Both are validated for T2VA only. +1. Ascend-friendly attention: the MindIE-SD `ascend_laser_attention` fused + kernel selected through `MINDIE_SD_FA_TYPE`, carried by the FLASH_ATTN + backend with a mask-free packed varlen path and K/V prefix slicing driven + by the packed `cu_seqlens` metadata. +2. Ascend-friendly fused operators: RMSNorm, AddRMSNorm, GQA, SwiGLU, and + rotary position embedding from MindIE-SD. +3. RainFusion block-sparse attention, optionally with EQBSA mixed INT8/FP8 + precision — see + [RainFusion](../../docs/user_guide/diffusion/attention_backends/rainfusion.md). +4. Online quantization: INT8, MXFP8, and MXFP4 — see + [online quantization](../../docs/user_guide/quantization/online.md). +5. Parallelism: USP, text-encoder TP, and native VAE `tile` patch + parallelism, plus distributed layerwise offload. +6. Post-training: few-step distillation via the FlashGen online LoRA. -**RainFusion block-sparse attention** — switch the attention backend, keeping -every other flag unchanged: +## Benchmarks (Atlas 800I A2 / A3) -```bash - --diffusion-attention-backend RAINFUSION_ATTN -``` +Measured with `--enable-diffusion-pipeline-profiler` on vLLM-Omni 0.28.0: +lossless rows use the recommended lossless configuration, lossy rows the +recommended lossy configuration. Memory figures are per device. -**INT8 online quantization** — add one flag to the server command above: +### Atlas 800I A2 -```bash - --quantization int8 -``` +| Task | Resolution | Frames | Duration (s) | Config | E2E (s) | DiT total (s) | DiT per-step (s) | DiT steps | Text encode (s) | Ref video encode (s) | Ref audio encode (s) | VAE decode (s) | Ref preprocess (s) | Post-process (s) | CPU MP4 (s) | Resident weights (GB) | Peak memory (GB) | +| ---- | ---- | ---- | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| t2va | 1344x768 | 124 | 5 | lossless | 150.06 | 139.9 | 2.85 | 49 | 0.83 | NA | NA | 6.76 | NA | 0.462 | 1.18 | 2.97 | 24.7 | +| t2va | 1344x768 | 362 | 15 | lossless | 786.24 | 761.36 | 15.53 | 49 | 1.20 | NA | NA | 15.87 | NA | 1.98 | 2.95 | 2.97 | 36.21 | +| t2va | 1344x768 | 124 | 5 | lossy | 108.1 | 95.85 | 1.95 | 49 | 0.9 | NA | NA | 6.73 | NA | 0.745 | 1.04 | 2.55 | 22.47 | +| t2va | 1344x768 | 362 | 15 | lossy | 552.51 | 526.04 | 10.73 | 49 | 1.28 | NA | NA | 16.07 | NA | 1.90 | 3.76 | 2.55 | 35.11 | +| ref2va | 1344x768 | 124 | 5 | lossless | 474.89 | 421.68 | 8.6 | 49 | 2.42 | 13.31 | 0.38 | 6.47 | 3.94 | 0.489 | 1.16 | 3.4 | 23.45 | +| ref2va | 1344x768 | 362 | 15 | lossless | 3414.18 | 3336.98 | 68.08 | 49 | 6.48 | 30.85 | 0.59 | 15.8 | 12.59 | 1.99 | 3.08 | 3.4 | 36.28 | +| ref2va | 1344x768 | 124 | 5 | lossy | 380.92 | 326.50 | 6.66 | 49 | 2.47 | 13.48 | 0.38 | 6.36 | 4.68 | 0.591 | 1.22 | 2.55 | 22.57 | +| ref2va | 1344x768 | 362 | 15 | lossy | 2332.04 | 2249.63 | 45.89 | 49 | 6.33 | 35.1 | 0.59 | 15.72 | 11.07 | 1.35 | 2.3 | 2.55 | 35.58 | -Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks -over the whole sequence, so ring parallelism would split away the keys it -needs. Scale with `--usp` instead. +### Atlas 800I A3 -## Memory and attention optimizations (A3) +| Task | Resolution | Frames | Duration (s) | Config | E2E (s) | DiT total (s) | DiT per-step (s) | DiT steps | Text encode (s) | Ref video encode (s) | Ref audio encode (s) | VAE decode (s) | Ref preprocess (s) | Post-process (s) | CPU MP4 (s) | Resident weights (GB) | Peak memory (GB) | +| ---- | ---- | ---- | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| t2va | 1344x768 | 124 | 5 | lossless | 114.19 | 105.87 | 2.16 | 49 | 0.48 | NA | NA | 6.56 | NA | 0.001 | 0.46 | 3.49 | 24.7 | +| t2va | 1344x768 | 362 | 15 | lossless | 599.29 | 580.38 | 11.84 | 49 | 1.83 | NA | NA | 14.95 | NA | 0.001 | 1.02 | 3.49 | 36.21 | +| t2va | 1344x768 | 124 | 5 | lossy | 84.82 | 77.47 | 1.58 | 49 | 0.57 | NA | NA | 6.23 | NA | 0.001 | 0.61 | 2.63 | 22.47 | +| t2va | 1344x768 | 362 | 15 | lossy | 401.43 | 383.56 | 7.82 | 49 | 1.72 | NA | NA | 14.83 | NA | 0.001 | 1.14 | 2.63 | 35.11 | +| ref2va | 1344x768 | 124 | 5 | lossless | 355.53 | 311.86 | 6.36 | 49 | 1.89 | 12.37 | 0.24 | 5.55 | 2.49 | 0.001 | 0.43 | 3.49 | 23.45 | +| ref2va | 1344x768 | 362 | 15 | lossless | 2339.56 | 2277.61 | 46.47 | 49 | 5.61 | 34.46 | 0.36 | 13.61 | 6.83 | 0.001 | 1.11 | 3.49 | 36.28 | +| ref2va | 1344x768 | 124 | 5 | lossy | 306.82 | 247.72 | 5.05 | 49 | 1.88 | 12.71 | 0.24 | 5.24 | 2.45 | 0.001 | 0.54 | 2.63 | 22.57 | +| ref2va | 1344x768 | 362 | 15 | lossy | 1695.26 | 1634.58 | 33.35 | 49 | 5.57 | 33.27 | 0.36 | 13.51 | 6.96 | 0.001 | 1.22 | 2.63 | 35.58 | -### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP - -Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the -combined MiniMax-H3 service at 768P. Enable **one** of the following at -server startup: - -**Distributed layerwise offload** — DiT layers are offloaded with parameters -gathered across the parallel group instead of replicated per rank: - -```bash - --enable-distributed-layerwise-offload -``` - -Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s -input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory -usage is high with this option. - -**HSDP** — hybrid sharded data parallelism for the DiT parameters. The -multi-stream memory-reuse knob is mandatory with this flag: - -```bash -export MULTI_STREAM_MEMORY_REUSE=2 -``` - -```bash - --use-hsdp -``` - -Host memory usage is small, but large shapes may still OOM; HBM usage -optimizations for this path are ongoing. - -### FLASH_ATTN backend with MindIE-SD - -Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see -[§ Environment](#environment)). On NPU this backend carries most of the -memory-reduction and performance work: a mask-free packed varlen path that -never materializes the quadratic `full_qk` padding mask, and K/V prefix -slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3 -transformer. - -### Optional: LaserAttention fused kernel - -For an additional attention speedup, select the Ascend LaserAttention fused -kernel before starting the server: - -```bash -export MINDIE_SD_FA_TYPE="ascend_laser_attention" -``` - -This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies -exact power-of-two input pre-scaling (`laser_input_scale=256`) so the -kernel's fp16 workspace cannot overflow on outlier activations. Measured -speedup numbers will be added here. - - -## HTTP API examples - -The request format is identical to the GPU recipe; see -[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples). -Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU. - -## Key parameters - -Same as the GPU recipe; see -[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters). -The validated resolution on NPU is 768P (e.g. 1344x768). - -## Validated NPU evidence - -Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1, -PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU -configuration above: - -| Workload | Configuration | -|----------|---------------| -| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile | -| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile | - -These measurements describe the validated shapes rather than a general -throughput guarantee. +These numbers describe the validated shapes rather than a general throughput +guarantee. ## Known limitations -- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads - both task-specific DiTs, and loads one copy of every shared component. +- Task serving is partitioned by `--task-type`: T2VA/FL2VA and Ref2VA load + different DiTs, so switching between them requires a restart with the other + partition. - H3 currently executes one generation request per diffusion batch. - The first regional-compile request is a warmup and should not be included in steady-state performance measurements. @@ -250,17 +380,19 @@ throughput guarantee. the same HTTP request contract. - VAE patch parallelism requires size 1 or the full DiT group size and supports the H3 native `tile` mode only. -- RainFusion block-sparse attention and INT8 quantization are validated for - T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA. -- Online INT8 quantization can be combined with distributed layerwise offload - while AllGather is enabled. Keep AllGather enabled for SP>1; use - `--dlo-no-use-allgather` only when collective communication is unavailable. +- The lossy configuration (RainFusion + INT8) is validated for T2VA and + Ref2VA; use the lossless configuration for FL2VA or re-validate first. ## Additional resources - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide -- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention) - — RainFusion knobs and tuning -- [Int8 quantization](../../docs/user_guide/quantization/int8.md) +- [MiniMax-H3-NPU-950PR.md](MiniMax-H3-NPU-950PR.md) — four-card + Ascend 950PR / 950DT guide +- [NPU installation](../../docs/getting_started/installation/npu.md) +- [RainFusion attention](../../docs/user_guide/diffusion/attention_backends/rainfusion.md) + — block-sparse knobs and tuning +- [Online quantization](../../docs/user_guide/quantization/online.md) + — INT8 / MXFP8 / MXFP4 +- [Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) - [Supported models](../../docs/models/supported_models.md) - [Video API](../../docs/serving/videos_api.md) From d0755b57169352c07416ad925385bdc907781e8c Mon Sep 17 00:00:00 2001 From: Chenchao Xu Date: Thu, 17 Sep 2026 20:18:51 +0800 Subject: [PATCH 31/99] =?UTF-8?q?[Model]=20Add=20=CF=800.5=20(Pi0.5)=20VLA?= =?UTF-8?q?=20model=20support=20(#6950)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: chenchaox --- .buildkite/cuda/test-nightly.yml | 11 + docs/models/supported_models.md | 1 + recipes/README.md | 3 + recipes/lerobot/Pi05.md | 154 +++ tests/diffusion/models/pi05/__init__.py | 2 + .../diffusion/models/pi05/test_pi05_parity.py | 500 ++++++++ .../diffusion/models/pi05/test_pi05_units.py | 1107 +++++++++++++++++ .../e2e/online_serving/test_pi05_expansion.py | 86 ++ .../openai_api/test_openpi_serving.py | 52 +- tests/helpers/runtime.py | 3 + tests/model_tests/diffusion/test_alignment.py | 1 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/pi05.yaml | 51 + vllm_omni/diffusion/data.py | 5 + vllm_omni/diffusion/models/pi05/__init__.py | 71 ++ vllm_omni/diffusion/models/pi05/config.py | 471 +++++++ .../diffusion/models/pi05/modeling_pi05.py | 898 +++++++++++++ .../diffusion/models/pi05/pipeline_pi05.py | 311 +++++ .../diffusion/models/pi05/processor_pi05.py | 568 +++++++++ .../diffusion/models/pi05_pipeline_config.py | 30 + vllm_omni/diffusion/registry.py | 7 + vllm_omni/entrypoints/openpi/serving.py | 14 +- 22 files changed, 4342 insertions(+), 6 deletions(-) create mode 100644 recipes/lerobot/Pi05.md create mode 100644 tests/diffusion/models/pi05/__init__.py create mode 100644 tests/diffusion/models/pi05/test_pi05_parity.py create mode 100644 tests/diffusion/models/pi05/test_pi05_units.py create mode 100644 tests/e2e/online_serving/test_pi05_expansion.py create mode 100644 vllm_omni/deploy/pi05.yaml create mode 100644 vllm_omni/diffusion/models/pi05/__init__.py create mode 100644 vllm_omni/diffusion/models/pi05/config.py create mode 100644 vllm_omni/diffusion/models/pi05/modeling_pi05.py create mode 100644 vllm_omni/diffusion/models/pi05/pipeline_pi05.py create mode 100644 vllm_omni/diffusion/models/pi05/processor_pi05.py create mode 100644 vllm_omni/diffusion/models/pi05_pipeline_config.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 1fdb1ef4a..adc94f9e9 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -223,6 +223,17 @@ steps: -m "full_model and diffusion and H100 and B200 and cards_1" --run-level "full_model" + - label: ":robot_face: Robot Policy OpenPI · π0.5 · H100 · Single-GPU" + timeout_in_minutes: 120 + commands: + - pip install --no-deps "openpi-client @ git+https://github.com/Physical-Intelligence/openpi.git@215abfb217dbac7d5f1273282331b9b1866c0479#subdirectory=packages/openpi-client" + - >- + pytest -sv + tests/e2e/online_serving/test_pi05_expansion.py + -m "full_model and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 42971ed97..f481154fb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -58,6 +58,7 @@ th { | `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) | +| `Pi05Pipeline` | π0.5 VLA | `lerobot/pi05_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi05.md) | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/recipes/README.md b/recipes/README.md index 078466e80..a7c10ca14 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -29,6 +29,8 @@ Example layout: ```text recipes/ + lerobot/ + Pi05.md Qwen/ Qwen3-Omni-H100.md Qwen3-TTS-A100.md @@ -53,6 +55,7 @@ recipes/ | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | | [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) | +| [`lerobot/Pi05.md`](./lerobot/Pi05.md) | Online OpenPI serving for the π0.5 vision-language-action policy | 1x H100 80GB | | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB | | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | diff --git a/recipes/lerobot/Pi05.md b/recipes/lerobot/Pi05.md new file mode 100644 index 000000000..af648f5ad --- /dev/null +++ b/recipes/lerobot/Pi05.md @@ -0,0 +1,154 @@ +# Pi0.5 (π0.5) VLA + +> Vision-Language-Action policy serving over the OpenPI realtime websocket + +## Summary + +- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot) +- Model: `lerobot/pi05_base` +- Task: Vision-Language-Action — multi-camera images + a language instruction + + robot proprioceptive state → a continuous action chunk + `[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0.5 + does **not** emit text tokens. +- Mode: Online serving via the OpenPI realtime robot API + (`/v1/realtime/robot/openpi` websocket) +- Maintainer: Community + +## When to use this recipe + +Use this recipe to deploy `lerobot/pi05_base` as a robot policy server: a robot +client connects over the OpenPI websocket and, per control step, sends an +observation (camera frames + state + instruction) and receives an action chunk. + +A single pipeline class (`Pi05Pipeline`) serves the policy. π0.5 is stateless +across calls (first-order Markov, no KV reuse), so the protocol's `session_id` / +`reset` are accepted but ignored. + +This recipe covers the functional path only. Serving runs in eager mode, with no +CUDA Graph capture, no Triton kernels and no prefix or encoder caching, and both +`float32` and `bfloat16` are supported. Two π0.5 capabilities are not +implemented, and a checkpoint that declares either is rejected when it loads: + +- **MEM** short-horizon observation memory (`use_visual_memory`, + `use_proprioceptive_memory`) +- **RTC** real-time chunking (`rtc_config`) + +A checkpoint with `use_relative_actions` must also carry `action_feature_names` +in `config.json`. + +## Weights + +| Component | Params | Part of | +| --- | ---: | --- | +| Vision tower, SigLIP So400m/14 | 412 M | PaliGemma | +| Language model, Gemma-2B, with the multimodal projector | 2.51 B | PaliGemma | +| Action expert, Gemma-300M, with AdaRMS and the action/time heads | 693 M | π0.5 | +| Total | 3.62 B | 13.5 GiB in float32 | + +| Artifact | Repository | +| --- | --- | +| `model.safetensors`, 812 tensors, ~14.5 GB | `lerobot/pi05_base` | +| `policy_preprocessor.json` + `policy_preprocessor_step_*.safetensors` | `lerobot/pi05_base` | +| Tokenizer, gated | `google/paligemma-3b-pt-224` | + +## Hardware Support + +### 1x H100 80GB (Online serving) + +The nightly end-to-end job runs on one H100, and serving was also verified on an +RTX 4090 and an RTX PRO 5000. + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM / vLLM-Omni: use the commit you are deploying from +- Hugging Face: `google/paligemma-3b-pt-224` license accepted, `hf auth login` +- OpenPI client extras (robot side, and for the e2e test): `openpi-client`, + `websockets`, `msgpack`, `msgpack-numpy`. CI pins `openpi-client` to OpenPI + revision `215abfb217dbac7d5f1273282331b9b1866c0479`. + +#### Command + +```bash +vllm serve lerobot/pi05_base \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --deploy-config vllm_omni/deploy/pi05.yaml \ + --enforce-eager --disable-log-stats +``` + +The deploy config (`vllm_omni/deploy/pi05.yaml`) declares a single diffusion +stage (`Pi05Pipeline`), `float32`, `max_num_seqs: 1`, `tokenizer_max_length: 200` +and `state_num_bins: 256`, together with the `policy_server_config` handshake +metadata (3 cameras, `joint_position`, action horizon 50, action dim 32) that +enables the OpenPI realtime serving layer. Keep `action_horizon`, `action_dim` +and `image_resolution` in that block in sync with `model_config`. + +To serve in bfloat16, set `dtype: bfloat16` in the same file. + +#### Verification + +The observation is a flat dict per inference: + +```python +{ + "observation.images.base_0_rgb": np.uint8[H, W, 3], # or float32 in [0, 1] + "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3], + "observation.images.right_wrist_0_rgb": np.uint8[H, W, 3], + "state": np.float32[state_dim], # must match the checkpoint's declared width + "prompt": "pick up the red block", + "session_id": "", # accepted but ignored (π0.5 is stateless) + "sampling_params": {"num_inference_steps": 4}, # optional engine knobs +} +``` + +Camera keys must match the server's `image_feature_keys`, which is the +checkpoint's `input_features` order. If your robot uses different camera names, +set `model_config.image_key_map` in `pi05.yaml` to map raw obs keys → feature +keys. Every request needs at least one configured camera and a non-null state. + +`state` must have the width `input_features["observation.state"]` declares, 32 +for `lerobot/pi05_base`. That checkpoint ships no normalization statistics, so +its values must already be normalized. + +PIL and uint8 images use the `[0, 255]` domain, while floating ndarray and +tensor images must be finite and in `[0, 1]`. Both are converted with +`x * 2 - 1`. + +Engine knobs live under `sampling_params`. It currently holds +`num_inference_steps`, which overrides the deploy-config denoising steps for a +single request. + +### Target-hardware validation + +Client E2E round trip over the OpenPI websocket: msgpack pack, send, engine, +denoise, unpack. Eager, batch 1, 10 denoising steps, p50 of 30 round +trips after 3 warmup. + +| GPU | dtype | Client E2E p50 (ms), 1 / 2 / 3 cameras | Device peak (GiB) | +| --- | --- | ---: | ---: | +| H100 80 GB HBM3 (sm_90) | float32 | 282 / 356 / 398 | 15.0 | +| H100 80 GB HBM3 (sm_90) | bfloat16 | 290 / 301 / 312 | 9.1 | +| RTX PRO 5000 72 GB (sm_120) | float32 | 185 / 231 / 281 | 14.8 | +| RTX PRO 5000 72 GB (sm_120) | bfloat16 | 156 / 177 / 203 | 8.9 | + +## Correctness + +π0.5's flow-matching kernel reproduces the LeRobot `PI05Policy` reference +exactly (`max|Δ| = 0`, CPU/float32, fixed noise, 1/2/3 cameras). The same file +also compares the prompt tokens carrying the discretized state, the +relative-action path, and a bfloat16 run against LeRobot in bfloat16: +`tests/diffusion/models/pi05/test_pi05_parity.py`. + +## References + +- Model blog: +- LeRobot π0.5 docs: +- Checkpoint: +- Pipeline: [`vllm_omni/diffusion/models/pi05/pipeline_pi05.py`](../../vllm_omni/diffusion/models/pi05/pipeline_pi05.py) +- Deploy config: [`vllm_omni/deploy/pi05.yaml`](../../vllm_omni/deploy/pi05.yaml) +- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi05/`](../../tests/diffusion/models/pi05/), + OpenPI websocket e2e in + [`tests/e2e/online_serving/test_pi05_expansion.py`](../../tests/e2e/online_serving/test_pi05_expansion.py) diff --git a/tests/diffusion/models/pi05/__init__.py b/tests/diffusion/models/pi05/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/diffusion/models/pi05/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/diffusion/models/pi05/test_pi05_parity.py b/tests/diffusion/models/pi05/test_pi05_parity.py new file mode 100644 index 000000000..fb30cb445 --- /dev/null +++ b/tests/diffusion/models/pi05/test_pi05_parity.py @@ -0,0 +1,500 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 parity against LeRobot. + +Reference: **lerobot 0.6.1**, transformers 5.5.4, torch 2.11.0+cpu. The oracle +is LeRobot's own code, so a later version changing its π0.5 implementation +fails this without anything here regressing — diff ``policies/pi05/`` first. + +Compares final action chunks under fixed noise: float32 exactly, bfloat16 +against LeRobot's own bfloat16 layout within ``BF16_ATOL``. + +Needs a separate ``lerobot[pi]`` venv (its dependencies conflict with +vllm-omni's) with the vllm-omni ``pi05`` package importable:: + + python -m pytest tests/diffusion/models/pi05/test_pi05_parity.py -v -s + +Skips itself when LeRobot is absent. ``PI05_PARITY_MODEL_PATH`` points at a +local checkpoint to skip the HF download. +""" + +from __future__ import annotations + +import copy +import importlib.util +import os + +import pytest +import torch + +# local_model: needs real weights + a lerobot venv, so it runs locally rather +# than in the ready-CI. Additionally gated on lerobot being importable. +_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None + +# cpu: the comparison runs on CPU in float32 so both implementations see +# identical numerics; local_model: it needs a lerobot venv that CI does not have. +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# ─── Config (fixed; matches LeRobot defaults for a π0.5 checkpoint) ──── +DEVICE = "cpu" +DTYPE_STR = "float32" +ATOL = 1e-4 +BF16_ATOL = 5e-2 # measured 2.95e-2, reproducible; margin is for version drift +NUM_STEPS = 10 +BATCH_SIZE = 2 +ACTION_DIM = 32 +STATE_DIM = 32 +ACTION_HORIZON = 50 +MAX_TOKEN_LEN = 200 # π0 uses 48 +NUM_STATE_BINS = 256 + +MODEL_PATH = os.environ.get("PI05_PARITY_MODEL_PATH", "lerobot/pi05_base") + +CAMERAS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb") + + +def _resolve_checkpoint_dir() -> str: + if os.path.isdir(MODEL_PATH): + return MODEL_PATH + from huggingface_hub import snapshot_download + + return snapshot_download(repo_id=MODEL_PATH, repo_type="model") + + +# ─── Dummy dataset stats ────────────────────────────────────────────── +def _dummy_dataset_stats() -> dict: + """Identity-ish stats. π0.5 normalizes with QUANTILES by default, so q01/q99 + carry the load here where π0's mean/std would. + + Using dummy stats means this test needs no checkpoint with real + ``norm_stats`` — it verifies numerical equivalence, not calibration. + """ + return { + "observation.state": { + "mean": torch.zeros(STATE_DIM), + "std": torch.ones(STATE_DIM), + "q01": -torch.ones(STATE_DIM), + "q99": torch.ones(STATE_DIM), + }, + "action": { + "mean": torch.zeros(ACTION_DIM), + "std": torch.ones(ACTION_DIM), + "q01": -torch.ones(ACTION_DIM), + "q99": torch.ones(ACTION_DIM), + }, + "images": { + cam: { + "mean": torch.zeros(3, 224, 224), + "std": torch.ones(3, 224, 224), + "q01": torch.zeros(3, 224, 224), + "q99": torch.ones(3, 224, 224), + } + for cam in CAMERAS + }, + } + + +def _create_dummy_batch(batch_size: int = BATCH_SIZE, num_views: int = 3, device: str = DEVICE) -> dict: + """Reproducible dummy inputs — identical across both implementations.""" + g = torch.Generator(device="cpu").manual_seed(0) + prompt = "Pick up the red block and place it in the bin" + batch = { + "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device), + "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device), + "task": [prompt for _ in range(batch_size)], + } + for cam in CAMERAS[:num_views]: + batch[f"observation.images.{cam}"] = torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to( + device + ) + return batch + + +# ─── Instantiation ──────────────────────────────────────────────────── +def _instantiate_lerobot(use_relative_actions: bool = False, action_feature_names=None, dtype: str = DTYPE_STR): + from lerobot.policies.pi05 import PI05Policy + from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors + + policy = PI05Policy.from_pretrained(MODEL_PATH, strict=True, dtype=dtype) + policy.to(DEVICE) + policy.config.device = DEVICE + policy.config.use_relative_actions = use_relative_actions + if action_feature_names is not None: + # Must land on the config *before* the processors are built: + # make_pi05_pre_post_processors captures action_feature_names into + # RelativeActionsProcessorStep at construction time, and its _build_mask + # treats ``action_names is None`` as "shift every dim" — including the + # gripper that relative_exclude_joints is supposed to hold back. Setting + # it afterwards leaves the step with the all-True mask and silently + # compares two different transforms. + policy.config.action_feature_names = action_feature_names + policy.eval() + + pre, post = make_pi05_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats()) + return policy, pre, post + + +def _instantiate_vllm_omni( + use_relative_actions: bool = False, + action_feature_names=None, + dtype: str = DTYPE_STR, +): + """Build the vllm-omni π0.5 model in isolation (no pipeline, no engine).""" + from vllm_omni.diffusion.models.pi05 import Pi05Config, Pi05ForActionPrediction + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype + + cfg = Pi05Config( + max_action_dim=ACTION_DIM, + max_state_dim=STATE_DIM, + chunk_size=ACTION_HORIZON, + num_inference_steps=NUM_STEPS, + tokenizer_max_length=MAX_TOKEN_LEN, + state_num_bins=NUM_STATE_BINS, + dtype=dtype, + use_relative_actions=use_relative_actions, + relative_exclude_joints=["gripper"] if action_feature_names else [], + action_feature_names=action_feature_names, + ) + model = Pi05ForActionPrediction(cfg) + _set_inference_dtype(model, getattr(torch, dtype)) + model.to(device=DEVICE) + _load_lerobot_weights(model) + model.eval() + return model, cfg + + +def _load_lerobot_weights(model): + import safetensors.torch + + path = os.path.join(_resolve_checkpoint_dir(), "model.safetensors") + state = safetensors.torch.load_file(path) + filled = model.load_weights(list(state.items())) + assert filled, "no weights were loaded — the remap rules are broken" + + +def _extract_lerobot_model_inputs(lerobot_policy, processed_batch): + """Mimic what ``PI05Policy.predict_action_chunk`` feeds into + ``self.model.sample_actions``. Using these *exact* tensors for vllm-omni + means any divergence must come from the core model, not preprocessing. + + Note π0.5 has no ``prepare_state`` step: the state is already inside + ``lang_tokens``. + """ + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + + images, img_masks = lerobot_policy._preprocess_images(processed_batch) + return images, img_masks, processed_batch[OBS_LANGUAGE_TOKENS], processed_batch[OBS_LANGUAGE_ATTENTION_MASK] + + +def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor: + g = torch.Generator(device="cpu").manual_seed(42) + return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device) + + +# ─── Main parity test ───────────────────────────────────────────────── +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize("num_views", [1, 2, 3]) +def test_pi05_vllm_omni_vs_lerobot(num_views): + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot() + omni_model, _ = _instantiate_vllm_omni() + + raw_batch = _create_dummy_batch(num_views=num_views) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed) + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + # The input contract: 256 tokens per view + a constant 200 text tokens. + assert lang_tokens.shape[1] == MAX_TOKEN_LEN, "π0.5 must pad text to exactly 200 tokens" + + with torch.no_grad(): + lerobot_actions = lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + omni_actions = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + diff = (lerobot_actions.float() - omni_actions.float()).abs() + print(f"[parity] views={num_views} |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}") + if not torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL): + _diagnose_divergence(lerobot_policy.model, omni_model, images, img_masks, lang_tokens, lang_masks, noise) + assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL), ( + f"actions differ beyond atol={ATOL}; max_diff={diff.max().item():.2e}" + ) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_bfloat16_matches_lerobot_bfloat16(): + """Both sides in bfloat16, so this compares implementations rather than + measuring bfloat16's drift from float32.""" + # LeRobot applies the layout in PaliGemmaWithExpertModel.__init__ from + # config.dtype, so the policy must be built as bfloat16 rather than cast after. + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot(dtype="bfloat16") + + omni_model, _ = _instantiate_vllm_omni(dtype="bfloat16") + + raw_batch = _create_dummy_batch(num_views=3) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed) + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + with torch.no_grad(): + lerobot_actions = lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + omni_actions = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + diff = (lerobot_actions.float() - omni_actions.float()).abs() + print(f"[parity] bfloat16 |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}") + assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=BF16_ATOL), ( + f"bfloat16 actions differ beyond atol={BF16_ATOL}; max_diff={diff.max().item():.2e}" + ) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_prompt_parity(): + """π0.5's state is carried as prompt tokens, so the tokenized prompt is + itself part of the contract. A bin-index mismatch here (e.g. from running + the discretizer before the normalizer) would otherwise only surface as a + small, plausible-looking action difference. + """ + from transformers import AutoTokenizer + + from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + apply_norm, + build_norm_stats, + build_pi05_prompt, + tokenize_prompt, + ) + + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot() + raw_batch = _create_dummy_batch(batch_size=1) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + + from lerobot.utils.constants import OBS_LANGUAGE_TOKENS + + lerobot_tokens = processed[OBS_LANGUAGE_TOKENS][0] + + stats = _dummy_dataset_stats()["observation.state"] + norm = build_norm_stats({"state": {"q01": stats["q01"], "q99": stats["q99"]}}, "state") + normalized = apply_norm(raw_batch["observation.state"][0].float(), norm) + prompt = build_pi05_prompt( + task=raw_batch["task"][0], + normalized_state=normalized.numpy(), + state_num_bins=NUM_STATE_BINS, + ) + tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224", padding_side="right") + ids, _ = tokenize_prompt(tokenizer, prompt, MAX_TOKEN_LEN) + + assert len(ids) == MAX_TOKEN_LEN + assert torch.equal(torch.tensor(ids), lerobot_tokens.cpu()), "π0.5 prompt tokens diverge from LeRobot" + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize("state_dim", [7, STATE_DIM]) +def test_pi05_prompt_parity_across_state_widths(state_dim): + """``pi05_base`` declares a 32-dim state, which is exactly ``max_state_dim``, + so the full-policy parity above cannot tell padding apart from the real + width. This drives LeRobot's tokenizer step directly, at both widths. + """ + from lerobot.lerobot_types import TransitionKey + from lerobot.policies.pi05.processor_pi05 import Pi05PrepareStateTokenizerProcessorStep + from lerobot.utils.constants import OBS_STATE + + from vllm_omni.diffusion.models.pi05.processor_pi05 import build_pi05_prompt + + g = torch.Generator().manual_seed(state_dim) + state = torch.randn(state_dim, generator=g, dtype=torch.float32) + task = "pick up the red block" + + step = Pi05PrepareStateTokenizerProcessorStep(max_state_dim=STATE_DIM) + transition = step( + { + TransitionKey.OBSERVATION: {OBS_STATE: state.unsqueeze(0)}, + TransitionKey.COMPLEMENTARY_DATA: {"task": [task]}, + } + ) + lerobot_prompt = transition[TransitionKey.COMPLEMENTARY_DATA]["task"][0] + + prompt = build_pi05_prompt(task=task, normalized_state=state.numpy()) + + assert prompt == lerobot_prompt + assert len(prompt.split("State: ")[1].split(";")[0].split()) == state_dim + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize( + "lerobot_mode,stats,ours,state,state_dim", + [ + ("MEAN_STD", {"mean": [0.0], "std": [1e-7]}, {"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1), + ("QUANTILES", {"q01": [0.0], "q99": [1e-7]}, {"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1), + ("QUANTILES", {"q01": [2.0], "q99": [2.0]}, {"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], 1), + ( + "QUANTILES", + {"q01": [0.0] * 4, "q99": [10.0] * 4}, + {"mode": "quantile", "q01": [0.0] * 4, "q99": [10.0] * 4}, + [-10.0, 0.0, 5.0, 20.0], + 4, + ), + ( + "MEAN_STD", + {"mean": [1.0, 2.0], "std": [0.5, 4.0]}, + {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 4.0]}, + [3.0, -6.0], + 2, + ), + ("MIN_MAX", {"min": [-1.0], "max": [3.0]}, {"mode": "min_max", "min": [-1.0], "max": [3.0]}, [2.0], 1), + ], +) +def test_pi05_state_normalization_parity(lerobot_mode, stats, ours, state, state_dim): + """State normalization against LeRobot's ``NormalizerProcessorStep``. + + The full-policy parity test feeds both sides stats with a wide range, so it + cannot see the denominator rule. A wrongly normalized state still discretizes + into a well-formed prompt: for ``q01=0, q99=1e-7`` a threshold substitution + moves the state token from bin 255 to bin 0. + """ + from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature + from lerobot.processor.normalize_processor import NormalizerProcessorStep + from lerobot.utils.constants import OBS_STATE + + from vllm_omni.diffusion.models.pi05.processor_pi05 import apply_norm, build_norm_stats + + step = NormalizerProcessorStep( + features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))}, + norm_map={FeatureType.STATE: getattr(NormalizationMode, lerobot_mode)}, + stats={OBS_STATE: {k: torch.tensor(v, dtype=torch.float32) for k, v in stats.items()}}, + ) + expected = step._apply_transform( + torch.tensor(state, dtype=torch.float32), OBS_STATE, FeatureType.STATE, inverse=False + ) + + actual = apply_norm(torch.tensor(state, dtype=torch.float32), build_norm_stats({"state": ours}, "state")) + + assert torch.allclose(actual, expected, atol=1e-6) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_relative_actions_parity(): + """With ``use_relative_actions=True`` the checkpoint's norm_stats live in + relative space, so the post-pipeline must add the raw state back. Compares + the full pre → model → post path against LeRobot's.""" + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + + from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05RelativeActions, + apply_norm, + build_norm_stats, + ) + + action_names = [f"joint_{i}" for i in range(ACTION_DIM - 1)] + ["gripper"] + lerobot_policy, lerobot_pre, lerobot_post = _instantiate_lerobot( + use_relative_actions=True, action_feature_names=action_names + ) + omni_model, cfg = _instantiate_vllm_omni(use_relative_actions=True, action_feature_names=action_names) + + raw_batch = _create_dummy_batch() + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks = lerobot_policy._preprocess_images(processed) + lang_tokens = processed[OBS_LANGUAGE_TOKENS] + lang_masks = processed[OBS_LANGUAGE_ATTENTION_MASK] + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + with torch.no_grad(): + lerobot_actions = lerobot_post( + lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + ) + omni_raw = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + rel = Pi05RelativeActions( + enabled=True, + exclude_joints=cfg.relative_exclude_joints, + action_names=action_names, + max_action_dim=ACTION_DIM, + ) + omni_actions = apply_norm(omni_raw, build_norm_stats(cfg.norm_stats, "action"), inverse=True) + # Per-sample reference state, matching LeRobot: its + # RelativeActionsProcessorStep caches the whole (B, D) state on the input + # side and shifts each sample by its own row. Passing a single row here + # would shift every sample by sample 0's state, which only agrees for B=1. + omni_actions = rel.to_absolute(omni_actions, raw_batch["observation.state"]) + + lerobot_actions = torch.as_tensor(lerobot_actions).float() + assert torch.allclose(lerobot_actions, omni_actions.float(), atol=ATOL), ( + "relative-action path diverges from LeRobot; " + f"max_diff={(lerobot_actions - omni_actions.float()).abs().max().item():.2e}" + ) + + +# ─── Per-stage divergence diagnostics ───────────────────────────────── +@torch.no_grad() +def _diagnose_divergence(lerobot_flow_model, omni_model, images, img_masks, lang_tokens, lang_masks, noise): + """Localize a numerical mismatch to a specific stage: + 1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang. + 2. Prefix KV cache layer 0 — PaliGemma LM attention. + 3. The AdaRMS timestep conditioning vector. + 4. A single denoise_step velocity at t=1.0 — action expert forward. + """ + from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + make_att_2d_masks, + prepare_attention_masks_4d, + ) + + lr_embs, lr_pad, lr_att = lerobot_flow_model.embed_prefix(images, img_masks, lang_tokens, lang_masks) + sg_embs, sg_pad, sg_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks) + print(f"[diag] prefix_embs max|Δ| = {(lr_embs.float() - sg_embs.float()).abs().max().item():.2e}") + print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_pad, sg_pad)}") + + img_len = 256 * len(images) + print(f"[diag] image slice max|Δ| = {(lr_embs[:, :img_len] - sg_embs[:, :img_len]).abs().max().item():.2e}") + print(f"[diag] lang slice max|Δ| = {(lr_embs[:, img_len:] - sg_embs[:, img_len:]).abs().max().item():.2e}") + + prefix_att_4d = prepare_attention_masks_4d(make_att_2d_masks(sg_pad, sg_att)) + prefix_pos = torch.cumsum(sg_pad, dim=1) - 1 + _, sg_kv = omni_model.paligemma_with_expert.forward( + attention_mask=prefix_att_4d, + position_ids=prefix_pos, + past_key_values=None, + inputs_embeds=[sg_embs, None], + use_cache=True, + ) + + t = torch.ones(noise.shape[0], dtype=torch.float32) + sg_cond = omni_model.embed_timestep(t) + print(f"[diag] adarms cond: shape={tuple(sg_cond.shape)} mean={sg_cond.mean().item():+.4f}") + + sg_vt = omni_model.denoise_step(sg_pad, sg_kv, noise, t) + print(f"[diag] denoise_step(t=1) v_t mean={sg_vt.mean().item():+.6f}") + + +if __name__ == "__main__": + test_pi05_vllm_omni_vs_lerobot(3) diff --git a/tests/diffusion/models/pi05/test_pi05_units.py b/tests/diffusion/models/pi05/test_pi05_units.py new file mode 100644 index 000000000..26671a59c --- /dev/null +++ b/tests/diffusion/models/pi05/test_pi05_units.py @@ -0,0 +1,1107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU unit tests for the π0.5 VLA model. + +Deliberately lightweight: config plumbing and the checkpoint-boundary rule, the +state discretization path, relative actions, and the weight-load remap. Runs on +CPU with no weights and no lerobot. The full LeRobot parity check lives in +``tests/diffusion/models/pi05/test_pi05_parity.py``. + +Most assertions here guard *silent* failure modes — cases where the wrong +behaviour still produces a well-shaped, finite action chunk: + + * normalization must run before state discretization; + * the normalization mode comes from the checkpoint, since π0.5 defaults to + quantile where π0 defaults to ``mean_std``; + * relative-action checkpoints must be honoured, and an unresolvable + ``relative_exclude_joints`` must raise rather than silently make every + dimension relative; + * MEM / RTC checkpoints must be rejected rather than served incorrectly. + + pytest tests/diffusion/models/pi05/test_pi05_units.py -v +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.models.pi05 import modeling_pi05 +from vllm_omni.diffusion.models.pi05.config import ( + Pi05Config, + UnsupportedCheckpointCapabilityError, + load_lerobot_norm_stats, + resolve_excluded_action_indices, +) +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + GemmaVariantConfig, + Pi05AdaRMSNorm, + Pi05ForActionPrediction, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, +) +from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05ImageProcessor, + Pi05Processor, + Pi05RelativeActions, + _assemble_model_inputs, + apply_norm, + as_state_vector, + build_norm_stats, + build_pi05_prompt, + discretize_state, + resize_with_pad, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ---------------------------------------------------------------------------- +# Pi05Config dataclass + checkpoint resolver +# ---------------------------------------------------------------------------- +_LEROBOT_CFG: dict[str, Any] = { + "type": "pi05", + "paligemma_variant": "gemma_2b", + "action_expert_variant": "gemma_300m", + "chunk_size": 50, + "max_action_dim": 32, + "max_state_dim": 32, + "num_inference_steps": 10, + "image_resolution": [224, 224], + "tokenizer_max_length": 200, + "dtype": "float32", + "input_features": { + "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.state": {"type": "STATE", "shape": [32]}, + }, + "output_features": {"action": {"type": "ACTION", "shape": [32]}}, +} + +_EXPECTED_CAMERA_ORDER = [ + "observation.images.base_0_rgb", + "observation.images.left_wrist_0_rgb", + "observation.images.right_wrist_0_rgb", +] + +_ACTION_NAMES = ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"] + + +def test_config_parses_lerobot_pi05_json(): + """``_LEROBOT_CFG`` mirrors the runtime keys of a real checkpoint.""" + c = Pi05Config.from_model_config(_LEROBOT_CFG) + assert c.tokenizer_max_length == 200 + assert c.state_num_bins == 256 + assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER + assert c.image_resolution == (224, 224) + + +def test_config_rejects_wrong_model_type(): + with pytest.raises(ValueError, match="type='pi05'"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, type="pi0")) + + +def test_config_rejects_unsupported_partial_action_chunk(): + with pytest.raises(UnsupportedCheckpointCapabilityError, match="n_action_steps"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, n_action_steps=10)) + + +@pytest.mark.parametrize( + "key,value", + [ + ("output_features", {"action": {"type": "ACTION", "shape": [33]}}), + ("input_features", {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [33]}}), + ], + ids=["action", "state"], +) +def test_config_rejects_a_schema_wider_than_the_model(key, value): + """``max_action_dim`` / ``max_state_dim`` are the widths the weights were + built for, so a wider checkpoint schema cannot be served.""" + with pytest.raises(ValueError, match="over max_state_dim"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + + +@pytest.mark.parametrize( + "key,value,attribute", + [ + ("output_features", {"action": {"type": "ACTION", "shape": [7]}}, "action_dim"), + ( + "input_features", + {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [7]}}, + "state_dim", + ), + ], + ids=["action", "state"], +) +def test_config_derives_the_real_width_from_the_schema(key, value, attribute): + """The declared width decides how many values reach the prompt and how many + columns come back, so a 7-joint checkpoint must not be padded out to 32.""" + config = Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + assert getattr(config, attribute) == 7 + + +def test_config_rejects_inconsistent_openpi_metadata(): + with pytest.raises(ValueError, match="policy_server_config.action_dim"): + Pi05Config.from_model_config( + dict( + _LEROBOT_CFG, + output_features={"action": {"type": "ACTION", "shape": [7]}}, + policy_server_config={"action_dim": 32}, + ) + ) + + +def test_config_non_square_resolution_raises(): + with pytest.raises(ValueError): + Pi05Config(image_resolution=(224, 256)) + + +@pytest.mark.parametrize("key", ["max_cameras", "num_inference_steps"]) +def test_config_rejects_counts_below_one(key): + """``build_model_inputs`` emits one image slot per camera and the Euler loop + steps ``-1 / num_inference_steps``; neither is defined at zero.""" + with pytest.raises(ValueError, match=f"{key} must be at least 1"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: 0})) + + +@pytest.mark.parametrize("key", ["state", "action"]) +def test_config_rejects_norm_stats_that_do_not_match_the_declared_width(key): + """A checkpoint's statistics and its feature schema come from different + files, so they can disagree.""" + stats = {key: {"mode": "quantile", "q01": [0.0] * 2, "q99": [1.0] * 2}} + with pytest.raises(ValueError, match=rf"norm_stats\[{key}\]\[q01\] has 2 entries, expected 32"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, norm_stats=stats)) + + +# ---------------------------------------------------------------------------- +# LeRobot normalization-stats sidecar +# ---------------------------------------------------------------------------- +# LeRobot keeps stats out of config.json: policy_preprocessor.json holds the +# structure and one safetensors file per stateful step holds the numbers. +def _normalized(raw, stats): + """The serving path's state normalization, without the pipeline.""" + return apply_norm(torch.as_tensor(raw, dtype=torch.float32), build_norm_stats({"state": stats}, "state")).numpy() + + +def _write_lerobot_checkpoint( + tmp_path, + *, + norm_map: dict[str, str] | None = None, + stats: dict[str, dict[str, list[float]]] | None = None, + write_state_file: bool = True, +) -> str: + """Write a minimal LeRobot-shaped checkpoint dir; return its path.""" + import json as _json + + import safetensors.torch as _st + + tmp_path.mkdir(parents=True, exist_ok=True) + state_file = "policy_preprocessor_step_3_normalizer_processor.safetensors" + step: dict = { + "registry_name": "normalizer_processor", + "config": { + "eps": 1e-08, + "features": {}, + "norm_map": norm_map if norm_map is not None else {"STATE": "QUANTILES", "ACTION": "QUANTILES"}, + }, + } + if write_state_file: + step["state_file"] = state_file + flat = { + f"{feature}.{stat}": torch.tensor(values, dtype=torch.float32) + for feature, entry in (stats or {}).items() + for stat, values in entry.items() + } + _st.save_file(flat, str(tmp_path / state_file)) + + (tmp_path / "policy_preprocessor.json").write_text( + _json.dumps({"name": "policy_preprocessor", "steps": [step]}), encoding="utf-8" + ) + (tmp_path / "config.json").write_text(_json.dumps(_config_for(stats)), encoding="utf-8") + return str(tmp_path) + + +def _config_for(stats: dict | None) -> dict: + """``_LEROBOT_CFG`` with the state/action widths the stats actually cover.""" + if not stats: + return _LEROBOT_CFG + config = dict(_LEROBOT_CFG) + state = next(iter(stats.get("observation.state", {}).values()), None) + action = next(iter(stats.get("action", {}).values()), None) + if state is not None: + features = dict(config["input_features"]) + features["observation.state"] = {"type": "STATE", "shape": [len(state)]} + config["input_features"] = features + if action is not None: + config["output_features"] = {"action": {"type": "ACTION", "shape": [len(action)]}} + return config + + +# A LeRobot state_dict carries every stat compute_stats emits, not just the two +# the norm_map selects. +_FULL_DATASET_STATS = { + "observation.state": { + "mean": [10.0, 20.0], + "std": [1.0, 2.0], + "min": [-5.0, -5.0], + "max": [5.0, 5.0], + "q01": [-1.0, -2.0], + "q99": [1.0, 2.0], + }, + "action": { + "mean": [0.5, 0.5], + "std": [0.1, 0.1], + "min": [-1.0, -1.0], + "max": [1.0, 1.0], + "q01": [-0.25, -0.5], + "q99": [0.25, 0.5], + }, +} + + +@pytest.mark.parametrize( + "declared,expected", + [ + ("QUANTILES", {"mode": "quantile", "q01": [-1.0, -2.0], "q99": [1.0, 2.0]}), + ("MEAN_STD", {"mode": "mean_std", "mean": [10.0, 20.0], "std": [1.0, 2.0]}), + ("IDENTITY", None), + ], +) +def test_load_lerobot_norm_stats_takes_the_mode_from_norm_map(tmp_path, declared, expected): + """The regression guard. ``_FULL_DATASET_STATS`` carries all six statistics, + as a real state_dict does, so an implementation that sniffed the present keys + would answer mean_std for every case here. Only norm_map distinguishes them. + """ + path = _write_lerobot_checkpoint( + tmp_path, norm_map={"STATE": declared, "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS + ) + stats = load_lerobot_norm_stats(path) + + if expected is None: + assert "state" not in stats + else: + assert stats["state"] == expected + assert stats["action"]["mode"] == "quantile" + + +@pytest.mark.parametrize("shape", ["no_sidecar", "no_state_file"]) +def test_load_lerobot_norm_stats_absent_stats_are_none(tmp_path, shape): + """``lerobot/pi05_base`` is the second shape: a normalizer step with no + stats attached. Neither is an error — the client is then expected to send an + already-normalized state.""" + if shape == "no_sidecar": + path = str(tmp_path) + else: + path = _write_lerobot_checkpoint(tmp_path, write_state_file=False) + + assert load_lerobot_norm_stats(path) is None + + +def test_load_lerobot_norm_stats_missing_declared_state_file_raises(tmp_path): + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + (tmp_path / "policy_preprocessor_step_3_normalizer_processor.safetensors").unlink() + + with pytest.raises(FileNotFoundError, match="normalizer state"): + load_lerobot_norm_stats(path) + + +def test_pipeline_rejects_missing_model_weights_before_initialization(tmp_path): + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.model_dir = str(tmp_path) + + with pytest.raises(FileNotFoundError, match="model.safetensors"): + pipeline._initialize_model() + + +def test_build_config_merges_the_deploy_yaml_over_the_checkpoint(tmp_path): + """The deploy yaml carries the runtime knobs and the checkpoint carries the + feature schema and the stats. Both have to land on one config.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.model_dir = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + + config = pipeline._build_config(SimpleNamespace(model_config={"type": "pi05", "num_inference_steps": 4})) + + assert config.num_inference_steps == 4 + assert (config.state_dim, config.action_dim) == (2, 2) + assert config.image_feature_keys == _EXPECTED_CAMERA_ORDER + assert config.norm_stats["state"]["mode"] == "quantile" + + +class _WideActionModel: + """Stands in for the 3.6B model, which is the one collaborator here that is + genuinely expensive to build. It returns the padded ``max_action_dim`` width + the real model emits, so the assertion is about the pipeline's cropping.""" + + def __init__(self, chunk_size: int, max_action_dim: int): + self._shape = (1, chunk_size, max_action_dim) + + def sample_actions(self, **kwargs): + del kwargs + return torch.ones(self._shape) + + +def test_pipeline_crops_actions_to_checkpoint_output_schema(monkeypatch): + from vllm_omni.diffusion.models.pi05 import pipeline_pi05 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + config = Pi05Config.from_model_config( + dict(_LEROBOT_CFG, output_features={"action": {"type": "ACTION", "shape": [7]}}) + ) + pipeline = object.__new__(pipeline_pi05.Pi05Pipeline) + pipeline.config = config + pipeline.tokenizer = object() + pipeline._device = torch.device("cpu") + pipeline.processor = Pi05Processor(config, object(), torch.device("cpu")) + pipeline.model = _WideActionModel(config.chunk_size, config.max_action_dim) + monkeypatch.setattr( + pipeline.processor, + "build_model_inputs", + lambda robot_obs: ([torch.empty(0)], [torch.empty(0)], torch.empty(0), torch.empty(0)), + ) + request = OmniDiffusionRequest( + prompt="", + sampling_params=OmniDiffusionSamplingParams( + extra_args={"robot_obs": {"state": np.zeros(config.max_state_dim)}}, + ), + request_id="crop-test", + ) + + result = pipeline.forward(request) + + assert result.output["actions"].shape == (config.chunk_size, 7) + + +def test_load_lerobot_norm_stats_unknown_mode_raises(tmp_path): + """Fail loud: an unreproducible mode served anyway is silently wrong.""" + path = _write_lerobot_checkpoint( + tmp_path, norm_map={"STATE": "SOMETHING_NEW", "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS + ) + with pytest.raises(ValueError, match="SOMETHING_NEW"): + load_lerobot_norm_stats(path) + + +def test_from_pretrained_backfills_sidecar_stats(tmp_path): + """End to end: the stats reach norm_stats without any deploy yaml.""" + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + config = Pi05Config.from_pretrained(path) + + assert config.norm_stats["state"]["mode"] == "quantile" + assert config.norm_stats["action"]["mode"] == "quantile" + + +def test_from_pretrained_does_not_override_explicit_norm_stats(tmp_path): + """config.json wins over the sidecar, matching LeRobot's override rule.""" + import json as _json + + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + explicit = {"state": {"mode": "min_max", "min": [0.0] * 32, "max": [1.0] * 32}} + (tmp_path / "config.json").write_text(_json.dumps({**_LEROBOT_CFG, "norm_stats": explicit}), encoding="utf-8") + + assert Pi05Config.from_pretrained(path).norm_stats == explicit + + +# ---------------------------------------------------------------------------- +# Checkpoint boundary: declared-but-unconsumed capabilities must raise +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize( + "key,value", + [ + ("use_visual_memory", True), + ("use_proprioceptive_memory", True), + ("rtc_config", {"mode": "trained", "inference_delay": 2}), + ("n_obs_steps", 6), + ], +) +def test_config_rejects_unsupported_capability(key, value): + """MEM and RTC change what a correct action chunk looks like and are not + visible in the weights. Serving such a checkpoint must fail loudly.""" + with pytest.raises(UnsupportedCheckpointCapabilityError, match=key): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + + +# ---------------------------------------------------------------------------- +# Relative actions +# ---------------------------------------------------------------------------- +def test_relative_actions_without_action_names_raises(): + """``action_feature_names`` is populated from dataset metadata at training + time. There is no dataset at serving time, so a relative-action checkpoint + that omits it cannot resolve ``relative_exclude_joints`` and must fail.""" + with pytest.raises(UnsupportedCheckpointCapabilityError, match="action_feature_names"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, action_feature_names=None)) + + +def test_relative_actions_empty_exclude_list_needs_no_action_names(): + c = Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, relative_exclude_joints=[])) + assert c.use_relative_actions + + +def test_resolve_excluded_action_indices_unresolvable_raises(): + with pytest.raises(UnsupportedCheckpointCapabilityError): + resolve_excluded_action_indices(["no_such_joint"], _ACTION_NAMES) + + +def _relative_step(enabled=True): + return Pi05RelativeActions( + enabled=enabled, + exclude_joints=["gripper"], + action_names=_ACTION_NAMES, + max_action_dim=32, + ) + + +def test_relative_mask_excludes_gripper_and_padding(): + rel = _relative_step() + assert rel.relative_mask[0] + assert not rel.relative_mask[6], "gripper must stay absolute" + assert not rel.relative_mask[len(_ACTION_NAMES) :].any(), "padded dims carry no signal" + assert rel.num_relative_dims == 6 + + +def test_relative_absolute_round_trip_is_exact(): + """The two directions are one object precisely so they cannot disagree.""" + rel = _relative_step() + g = torch.Generator().manual_seed(0) + actions = torch.randn(1, 50, 32, generator=g) + state = torch.randn(32, generator=g) + assert torch.allclose(rel.to_absolute(rel.to_relative(actions, state), state), actions, atol=1e-6) + + +def test_relative_rejects_mismatched_action_dim(): + rel = _relative_step() + with pytest.raises(ValueError, match="max_action_dim"): + rel.to_absolute(torch.zeros(1, 4, 8), torch.zeros(32)) + + +# ---------------------------------------------------------------------------- +# State discretization — π0.5's defining input path +# ---------------------------------------------------------------------------- +def test_discretize_underflows_to_minus_one_and_saturates_at_the_top(): + """Out-of-range input is asymmetric, and that asymmetry is LeRobot's. + + ``np.digitize`` returns 0 for anything below the first bin edge, so the + ``- 1`` makes a below-range state land in bin ``-1``; above-range input + saturates at ``num_bins - 1`` on its own. The checkpoint was trained with + ``" -1"`` in the state prompt for under-range dims, so clipping that to 0 + would change the tokens the model sees — LeRobot parity catches it. + """ + out = discretize_state(np.array([-9.0, 9.0], dtype=np.float32), num_bins=256) + assert out.tolist() == [-1, 255] + + +def test_normalize_state_does_not_clip_out_of_range(): + """The parity-critical contract: normalization is affine and nothing else. + + Clamping to [-1, 1] here would hide an under-range dimension from + ``discretize_state``, turning LeRobot's ``-1`` bin into ``0`` and changing + the state tokens in the prompt. LeRobot applies + ``2.0 * (x - q01) / (q99 - q01) - 1.0`` with no clip. + """ + stats = {"q01": [0.0] * 32, "q99": [10.0] * 32} + out = _normalized([-10.0, 20.0] + [5.0] * 30, stats) + assert out[0] == pytest.approx(-3.0) # would be -1.0 if clipped + assert out[1] == pytest.approx(3.0) # would be 1.0 if clipped + # and the under-range dim must reach the -1 bin end to end + assert discretize_state(out, num_bins=256)[0] == -1 + + +def test_prompt_serializes_the_declared_state_width(): + """LeRobot's tokenizer step discretizes the state at its real width and + never pads to ``max_state_dim``.""" + prompt = build_pi05_prompt(task="x", normalized_state=np.zeros(7, dtype=np.float32)) + bins = prompt.split("State: ")[1].split(";")[0].split() + assert len(bins) == 7 + + +def test_prompt_template_matches_lerobot(): + prompt = build_pi05_prompt( + task=" pick_up the red\nblock ", + normalized_state=np.zeros(4, dtype=np.float32), + ) + assert prompt.startswith("Task: pick up the red block, State: ") + assert prompt.endswith(";\nAction: ") + + +# ---------------------------------------------------------------------------- +# Input contract: fixed camera slots x 256 tokens + constant 200 text tokens +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("views,expected", [(1, 456), (2, 712), (3, 968)]) +def test_prefix_length_is_fixed_by_max_cameras(views, expected): + """``max_cameras × 256`` image tokens plus a constant ``tokenizer_max_length``. + + The text segment is padded to its maximum and the camera slots are padded to + ``max_cameras``, so the prefix shape does not vary with the request — only + the attention mask does. + """ + keys = _EXPECTED_CAMERA_ORDER[:views] + config = Pi05Config(max_cameras=views, tokenizer_max_length=200, image_feature_keys=keys) + observation = { + "state": np.zeros(32, dtype=np.float32), + "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in keys}, + } + + images, masks, lang_tokens, _ = _assemble_model_inputs( + observation, config, _FakeTokenizer(), torch.device("cpu"), None + ) + + assert len(images) == views + assert all(bool(mask) for mask in masks) + assert lang_tokens.shape[-1] == 200 + assert 256 * len(images) + lang_tokens.shape[-1] == expected + + +# ---------------------------------------------------------------------------- +# Shared math helpers (same contracts as π0) +# ---------------------------------------------------------------------------- +def test_gemma_variant_dims(): + vlm = get_gemma_config("gemma_2b") + assert (vlm.width, vlm.depth, vlm.mlp_dim, vlm.num_kv_heads, vlm.head_dim) == (2048, 18, 16384, 1, 256) + expert = get_gemma_config("gemma_300m") + assert (expert.width, expert.depth, expert.mlp_dim) == (1024, 18, 4096) + + +def test_sinusoidal_embedding_shape_and_oddity(): + out = create_sinusoidal_pos_embedding(torch.tensor([0.0, 1.0]), 64) + assert out.shape == (2, 64) + with pytest.raises(ValueError): + create_sinusoidal_pos_embedding(torch.tensor([0.0]), 63) + + +def test_make_att_2d_masks_respects_padding(): + pad = torch.tensor([[True, True, False]]) + att = torch.tensor([[0, 0, 0]]) + out = make_att_2d_masks(pad, att) + assert not out[0, :, 2].any() + + +def test_build_norm_stats_rejects_unknown_mode(): + with pytest.raises(ValueError, match="Unsupported normalization mode"): + build_norm_stats({"action": {"mode": "qunatile"}}, "action") + + +@pytest.mark.parametrize( + "entry", + [ + {"mode": "mean_std", "mean": [0.0]}, + {"mode": "min_max", "min": [0.0]}, + {"mode": "quantile", "q01": [0.0]}, + ], +) +def test_build_norm_stats_rejects_incomplete_stats(entry): + with pytest.raises(ValueError, match="requires"): + build_norm_stats({"action": entry}, "action") + + +def test_normalize_quantile_round_trip_and_padded_tail(): + stats = build_norm_stats({"action": {"q01": [0.0, 0.0], "q99": [2.0, 4.0]}}, "action") + x = torch.tensor([[0.5, -0.5, 7.0]]) # third entry is padding beyond the stats + back = apply_norm(apply_norm(x, stats), stats, inverse=True) + assert torch.allclose(back, x, atol=1e-5) + assert back[0, 2] == 7.0, "padded tail must pass through untouched" + + +# ---------------------------------------------------------------------------- +# Image processor +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("value,expected", [(0.0, -1.0), (0.25, -0.5), (1.0, 1.0)]) +@pytest.mark.parametrize("container", ["numpy", "torch"]) +def test_float_image_uses_lerobot_zero_one_domain(value, expected, container): + if container == "numpy": + image = np.full((4, 4, 3), value, dtype=np.float32) + else: + image = torch.full((3, 4, 4), value, dtype=torch.float32) + + out = Pi05ImageProcessor(image_size=4).preprocess_single(image) + + assert out.shape == (1, 3, 4, 4) + assert torch.allclose(out, torch.full_like(out, expected)) + + +@pytest.mark.parametrize("bad", [-0.01, 1.01, np.nan, np.inf]) +@pytest.mark.parametrize("container", ["numpy", "torch"]) +def test_float_image_rejects_values_outside_lerobot_domain(bad, container): + if container == "numpy": + image = np.full((4, 4, 3), bad, dtype=np.float32) + else: + image = torch.full((3, 4, 4), bad, dtype=torch.float32) + + with pytest.raises(ValueError): + Pi05ImageProcessor(image_size=4).preprocess_single(image) + + +def test_resize_with_pad_pads_with_minus_one(): + out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224) + assert out.shape == (1, 3, 224, 224) + assert out[0, 0, 0, 0] == -1.0 + + +class _FakeTokenizer: + def __call__(self, text, **kwargs): + del text + length = kwargs["max_length"] + return {"input_ids": [0] * length, "attention_mask": [1] * length} + + +@pytest.mark.parametrize("explicit_none", [False, True]) +def test_build_model_inputs_preserves_missing_middle_camera_slot(explicit_none): + config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER, max_cameras=3) + images = { + _EXPECTED_CAMERA_ORDER[0]: np.zeros((4, 4, 3), dtype=np.uint8), + _EXPECTED_CAMERA_ORDER[2]: np.full((4, 4, 3), 255, dtype=np.uint8), + } + if explicit_none: + images[_EXPECTED_CAMERA_ORDER[1]] = None + observation = { + "images": images, + "state": np.zeros(32, dtype=np.float32), + } + + slots, masks, _, _ = _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None) + + assert [bool(mask.item()) for mask in masks] == [True, False, True] + torch.testing.assert_close(slots[0], torch.full_like(slots[0], -1.0)) + torch.testing.assert_close(slots[1], torch.full_like(slots[1], -1.0)) + torch.testing.assert_close(slots[2], torch.full_like(slots[2], 1.0)) + + +def test_build_model_inputs_requires_configured_camera(): + config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER) + observation = { + "unknown_camera": np.zeros((4, 4, 3), dtype=np.uint8), + "state": np.zeros(32, dtype=np.float32), + } + with pytest.raises(ValueError, match="configured camera"): + _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None) + + +# ---------------------------------------------------------------------------- +# Model behaviour + weight-load remap (tiny action dims; full Gemma backbone) +# ---------------------------------------------------------------------------- +# Module *structure* (no ``state_proj``, ``time_mlp_{in,out}`` widths, AdaRMS on +# the expert norms but not the prefix) is not asserted here: ``load_weights`` +# audits every model parameter against the checkpoint and raises on a missing, +# unmatched or π0-shaped key, so a structural drift fails loudly at load time +# rather than silently. ``test_pi0_shaped_keys_are_not_silently_loaded`` guards +# that audit, and ``test_pi05_parity.py`` is the numerical oracle. +# Both Gemma towers keep their real widths but drop from 18 layers to 2 for +# every test in this section. Nothing here depends on depth: the weight remap +# matches by name, and the AR mask, prefix length and ODE determinism are +# shape-driven. Full-scale numerical correctness is ``test_pi05_parity.py``'s +# job, against the real checkpoint. +# +# Two dimensions must NOT be shrunk, both because something outside these +# configs is sized against them: +# +# * ``width`` — the vision tower's ``multi_modal_projector`` projects image +# embeddings to the PaliGemma text width so the two can be concatenated. +# Narrowing the text tower alone makes that concat fail. +# * the vision tower itself — So400m/14 on 224x224 is what yields exactly 256 +# tokens per view, which ``test_prefix_length_and_valid_len_for_each_view_count`` +# exists to verify. +# +# Depths are equal across the two towers on purpose: the action expert runs +# layer-by-layer alongside the VLM, so a mismatch would not be a valid model. +_SHRUNK_GEMMA = { + "gemma_2b": GemmaVariantConfig(2048, 2, 1024, 8, 1, 256), + "gemma_300m": GemmaVariantConfig(1024, 2, 512, 8, 1, 256), +} + + +@pytest.fixture(scope="module", autouse=True) +def _shrink_gemma_backbones(): + """Patch the module-level lookup ``Pi05ForActionPrediction.__init__`` calls.""" + patch = pytest.MonkeyPatch() + patch.setattr(modeling_pi05, "get_gemma_config", lambda variant: _SHRUNK_GEMMA[variant]) + yield + patch.undo() + + +def _tiny_pi05_model(): + """A fresh model. Use for tests that load weights or otherwise mutate it.""" + return Pi05ForActionPrediction(Pi05Config(max_action_dim=8, max_state_dim=8, chunk_size=4, n_action_steps=4)) + + +@pytest.fixture(scope="module") +def tiny_model(): + """Shared read-only model, so the tests that only inspect it pay for one + construction between them rather than one each.""" + return _tiny_pi05_model() + + +@pytest.fixture(autouse=True) +def _silence_pi05_loader(): + logging.getLogger("vllm_omni.diffusion.models.pi05.modeling_pi05").setLevel(logging.CRITICAL) + yield + + +@pytest.mark.slow +def test_suffix_has_no_state_token(tiny_model): + """π0's suffix is ``[state, actions...]`` with AR mask ``[1, 1, 0...]``; + π0.5's is actions only with ``[1, 0...]``.""" + embs, pad, att, cond = tiny_model.embed_suffix(torch.randn(1, 4, 8), torch.tensor([1.0])) + assert embs.shape[1] == 4 + assert att[0].tolist() == [1, 0, 0, 0] + assert cond.shape == (1, tiny_model.expert_width) + + +# Not marked slow: constructs one small norm layer, no backbone. +def test_adarms_norm_returns_gate_only_when_conditioned(): + norm = Pi05AdaRMSNorm(8, cond_dim=4) + out, gate = norm(torch.randn(1, 3, 8), torch.randn(1, 4)) + assert gate is not None and gate.shape == (1, 1, 8) + + plain = Pi05AdaRMSNorm(8, cond_dim=None) + out, gate = plain(torch.randn(1, 3, 8)) + assert gate is None + + +@pytest.mark.slow +def test_lm_head_remap_to_embed_tokens(): + model = _tiny_pi05_model() + target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + payload = torch.full(dict(model.named_parameters())[target].shape, 0.5) + model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)], strict=False) + assert torch.allclose(dict(model.named_parameters())[target], payload) + + +@pytest.mark.slow +def test_pi0_shaped_keys_are_not_silently_loaded(): + """A π0 checkpoint pointed at the π0.5 class would leave the AdaRMS expert + randomly initialized. ``state_proj`` is the tell.""" + model = _tiny_pi05_model() + with pytest.raises(RuntimeError, match="π0-shaped"): + model.load_weights([("model.state_proj.weight", torch.zeros(1024, 8))]) + + +@pytest.mark.slow +@pytest.mark.parametrize("num_real_views", [1, 2, 3]) +def test_prefix_length_is_fixed_and_valid_len_tracks_masks(num_real_views, tiny_model): + """The deployed model has three slots; only their validity varies.""" + live_text = 120 + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([index < num_real_views]) for index in range(3)] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.zeros(1, 200, dtype=torch.bool) + lang_mask[:, :live_text] = True + + embs, pad_masks, _ = tiny_model.embed_prefix(images, masks, lang, lang_mask) + assert embs.shape[1] == 256 * 3 + 200 + assert int(pad_masks.sum()) == 256 * num_real_views + live_text + + +@pytest.mark.slow +def test_sample_actions_shape_and_determinism(tiny_model): + """Flow matching is an ODE: fixed noise must give a bit-identical chunk.""" + model = tiny_model.eval() + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(42)) + + with torch.no_grad(): + a1 = model.sample_actions( + images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2 + ) + a2 = model.sample_actions( + images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2 + ) + assert a1.shape == (1, 4, 8) + assert torch.isfinite(a1).all() + assert torch.equal(a1, a2) + + +@pytest.mark.slow +def test_sample_actions_uses_request_generator(tiny_model): + model = tiny_model.eval() + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + + def run(seed): + return model.sample_actions( + images=images, + image_masks=masks, + lang_tokens=lang, + lang_masks=lang_mask, + num_steps=2, + generator=torch.Generator().manual_seed(seed), + ) + + with torch.no_grad(): + first = run(42) + repeated = run(42) + different = run(43) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different) + + +# ---------------------------------------------------------------------------- +# Serving dtype +# ---------------------------------------------------------------------------- +# The check that decides what runs lives on the pipeline, not on Pi05Config: +# Pi05Config.dtype records what the checkpoint declares and never reaches a +# cast, while _initialize_model casts using the top-level OmniDiffusionConfig. +@pytest.mark.parametrize( + "declared,expected", + [ + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + ], +) +def test_resolve_dtype_accepts_supported_serving_dtypes(declared, expected): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + assert Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared)) is expected + + +@pytest.mark.parametrize("declared", ["float16", "fp16", torch.float16]) +def test_resolve_dtype_rejects_unvalidated_dtypes(declared): + """float16 resolves to a real torch dtype, so a permissive lookup would run + the model in a precision nothing here has been checked in. + + Only dtypes that survive ``OmniDiffusionConfig`` normalization are worth + asserting on: it maps a string it does not recognize onto bfloat16 with a + warning, so e.g. ``"float64"`` never reaches this guard as float64. + """ + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + with pytest.raises(ValueError, match="Unsupported π0.5 dtype"): + Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared)) + + +def test_bfloat16_layout_is_applied_before_checkpoint_load(monkeypatch, tmp_path): + """Checkpoint tensors must be copied directly into their final dtypes. + + LeRobot casts only the inner PaliGemma/expert module; the outer action and + time projections remain float32. The layout is established before loading + so float32-retained weights never pass through bfloat16. + """ + from vllm_omni.diffusion.models.pi05 import pipeline_pi05 + + class _TinyPi05(torch.nn.Module): + def __init__(self): + super().__init__() + self.paligemma_with_expert = torch.nn.Module() + self.paligemma_with_expert.backbone = torch.nn.Linear(2, 2) + self.paligemma_with_expert.vision_tower = torch.nn.Linear(2, 2) + self.paligemma_with_expert.input_layernorm = torch.nn.LayerNorm(2) + self.paligemma_with_expert.register_buffer("floating_buffer", torch.ones(1)) + self.action_in_proj = torch.nn.Linear(2, 2) + self.action_out_proj = torch.nn.Linear(2, 2) + self.time_mlp_in = torch.nn.Linear(2, 2) + self.time_mlp_out = torch.nn.Linear(2, 2) + + # DiffusersLoader constructs custom pipelines under the requested default + # dtype. Start fully in bfloat16 to verify that the helper restores every + # float32-retained tensor rather than relying on the construction default. + model = _TinyPi05().to(dtype=torch.bfloat16) + observed_dtypes = {} + + def fake_load_checkpoint(_self, model_to_load): + observed_dtypes.update({name: param.dtype for name, param in model_to_load.named_parameters()}) + + (tmp_path / "model.safetensors").touch() + pipeline = object.__new__(pipeline_pi05.Pi05Pipeline) + pipeline.model_dir = str(tmp_path) + pipeline.config = object() + pipeline._device = torch.device("cpu") + pipeline._torch_dtype = torch.bfloat16 + monkeypatch.setattr(pipeline_pi05, "Pi05ForActionPrediction", lambda _config: model) + monkeypatch.setattr(pipeline_pi05.Pi05Pipeline, "_load_checkpoint", fake_load_checkpoint) + + initialized = pipeline._initialize_model() + + assert observed_dtypes["paligemma_with_expert.backbone.weight"] is torch.bfloat16 + assert model.paligemma_with_expert.floating_buffer.dtype is torch.bfloat16 + for name in ( + "paligemma_with_expert.vision_tower.weight", + "paligemma_with_expert.input_layernorm.weight", + "action_in_proj.weight", + "action_out_proj.weight", + "time_mlp_in.weight", + "time_mlp_out.weight", + ): + assert observed_dtypes[name] is torch.float32 + assert initialized is model + assert not initialized.training + + +@pytest.mark.slow +def test_bfloat16_runs_and_tracks_float32(tiny_model): + """bfloat16 is a supported serving dtype, so it has to produce a usable + chunk — and one that still resembles float32. + + The bound is deliberately loose. This is not a parity check (bfloat16 + deviates from float32 by ~1-2% of the action range on the real checkpoint, + which is why recipes/lerobot/Pi05.md tells operators to validate it on their + own task). What it catches is gross breakage: a cast that silently did not + happen, activations overflowing to inf, or a dtype mismatch that would make + the two outputs unrelated. + """ + import copy + + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(7)) + + def run(model, dtype): + with torch.no_grad(): + return model.sample_actions( + images=[img.to(dtype) for img in images], + image_masks=masks, + lang_tokens=lang, + lang_masks=lang_mask, + noise=noise.to(dtype), + num_steps=2, + ).float() + + reference = run(tiny_model.eval(), torch.float32) + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype + + bf16_model = copy.deepcopy(tiny_model).eval() + _set_inference_dtype(bf16_model, torch.bfloat16) + + # The cast actually happened — otherwise everything below passes trivially. + assert any(param.dtype is torch.bfloat16 for param in bf16_model.parameters()) + assert bf16_model.action_in_proj.weight.dtype is torch.float32 + assert bf16_model.action_out_proj.weight.dtype is torch.float32 + assert bf16_model.time_mlp_in.weight.dtype is torch.float32 + assert bf16_model.time_mlp_out.weight.dtype is torch.float32 + + actual = run(bf16_model, torch.bfloat16) + assert actual.shape == reference.shape + assert torch.isfinite(actual).all() + + scale = float(reference.abs().max()) + assert float((actual - reference).abs().max()) < max(0.5 * scale, 1e-2) + + +# ---------------------------------------------------------------------------- +# State normalization runs π0's stats path +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize( + "stats,state,expected", + [ + # A denominator that is small but not zero. ``normalize`` divides by + # ``std + eps``; replacing a small std with 1.0 would be a different map. + ({"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1e-7 / (1e-7 + 1e-8)), + ({"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1.0), + ({"mode": "quantile", "q01": [0.0], "q99": [1e-9]}, [1e-9], 1.0), + # An exactly zero range maps the bound to -1 instead of dividing by zero. + ({"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], -1.0), + ], +) +def test_normalize_state_uses_the_shared_denominator_rule(stats, state, expected): + """State normalization and action unnormalization are the same affine map, + so a narrow range keeps producing LeRobot's answer.""" + assert _normalized(state, stats)[0] == pytest.approx(expected, rel=1e-5) + + +# ---------------------------------------------------------------------------- +# State wire contract +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_state_rejects_non_finite_values(bad): + """``np.digitize`` maps nan and inf to the top bin and -inf to -1, all legal + state tokens, so a dead sensor would reach the model as an ordinary prompt.""" + with pytest.raises(ValueError, match="NaN or Inf"): + as_state_vector([0.0, bad], 2) + + +def test_state_rejects_a_multi_row_matrix(): + """Flattening would read a (2, 16) request as a 32-value state.""" + with pytest.raises(ValueError, match=r"\(D,\) or \(1, D\)"): + as_state_vector(np.zeros((2, 16), dtype=np.float32), 32) + + +# ---------------------------------------------------------------------------- +# Per-request denoising steps reach the model +# ---------------------------------------------------------------------------- +class _SpyModel: + """Records the ``num_steps`` ``forward`` hands to ``sample_actions``.""" + + def __init__(self): + self.seen: list = [] + self.generators: list = [] + + def sample_actions(self, *, images, image_masks, lang_tokens, lang_masks, num_steps, generator): + del images, image_masks, lang_tokens, lang_masks + self.seen.append(num_steps) + self.generators.append(generator) + return torch.zeros(1, 50, 32) + + +def _spy_pipeline(spy: _SpyModel): + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.config = Pi05Config.from_model_config(_LEROBOT_CFG) + pipeline.tokenizer = _FakeTokenizer() + pipeline._device = torch.device("cpu") + pipeline.model = spy + pipeline.processor = Pi05Processor(pipeline.config, _FakeTokenizer(), torch.device("cpu")) + return pipeline + + +def _spy_request(num_inference_steps, generator=None): + params = type("_Params", (), {})() + params.extra_args = { + "robot_obs": { + "state": np.zeros(32, dtype=np.float32), + "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in _EXPECTED_CAMERA_ORDER}, + } + } + params.num_inference_steps = num_inference_steps + params.generator = generator + request = type("_Request", (), {})() + request.sampling_params = params + return request + + +@pytest.mark.parametrize("requested", [None, 2, 25]) +def test_request_denoising_steps_reach_sample_actions(requested): + """The OpenPI E2E cannot see this: a pipeline that ignored the override and + ran the static default would still return a finite chunk of the right shape.""" + spy = _SpyModel() + output = _spy_pipeline(spy).forward(_spy_request(requested)) + + assert spy.seen == [requested] + assert output.output["actions"].shape == (50, 32) + + +def test_request_generator_reaches_sample_actions(): + generator = torch.Generator().manual_seed(42) + spy = _SpyModel() + + _spy_pipeline(spy).forward(_spy_request(2, generator)) + + assert spy.generators == [generator] + + +@pytest.mark.parametrize("bad", [0, -1, 2.5, True, "4"]) +def test_request_denoising_steps_must_be_a_positive_integer(bad): + spy = _SpyModel() + with pytest.raises(ValueError, match="num_inference_steps must be a positive integer"): + _spy_pipeline(spy).forward(_spy_request(bad)) + assert spy.seen == [] diff --git a/tests/e2e/online_serving/test_pi05_expansion.py b/tests/e2e/online_serving/test_pi05_expansion.py new file mode 100644 index 000000000..0c6fbfd34 --- /dev/null +++ b/tests/e2e/online_serving/test_pi05_expansion.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""E2E online serving test for π0.5 OpenPI websocket serving. + +Boots ``vllm serve --omni --deploy-config pi05.yaml`` and drives the real OpenPI +websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses +(handshake metadata → send observation → receive action chunk). Mirrors +``tests/e2e/online_serving/test_pi0_expansion.py``. Needs one H100 and the full +π0.5 checkpoint. + +The ``pi0_openpi_*`` helpers are protocol-level (handshake, send obs, receive +chunk) and not π0-specific, so π0.5 reuses them; only the deploy config and the +expected metadata differ. This file is selected explicitly by the nightly H100 +robot-policy job. + +The in-process LeRobot parity oracle lives separately in +``tests/diffusion/models/pi05/test_pi05_parity.py``. +""" + +from __future__ import annotations + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import ( + OmniServerParams, + get_open_port, + pi0_openpi_require_dependencies, + pi0_openpi_run_policy_session, + pi0_openpi_validate_session_result, +) + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +MODEL = "lerobot/pi05_base" + + +test_params = [ + OmniServerParams( + model=MODEL, + port=8093, + server_args=[ + "--deploy-config", + "vllm_omni/deploy/pi05.yaml", + "--enforce-eager", + "--disable-log-stats", + ], + env_dict={ + "ATTENTION_BACKEND": "torch", + "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA", + "VLLM_DISABLE_COMPILE_CACHE": "1", + "MASTER_PORT": str(get_open_port()), + }, + ) +] + + +@pytest.mark.full_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_pi05_openpi_online(omni_server): + pi0_openpi_require_dependencies() + + result = pi0_openpi_run_policy_session( + host=omni_server.host, + port=omni_server.port, + prompt="pick up the red block and place it in the bin", + session_id="pi05-online-e2e", + num_steps=2, + num_inference_steps=2, + ) + + # Asserts every returned chunk is [50, 32] + finite, and the handshake + # metadata matches pi05.yaml's policy_server_config. + pi0_openpi_validate_session_result(result) + + metadata = result["metadata"] + assert tuple(metadata["image_resolution"]) == (224, 224) + assert metadata["needs_wrist_camera"] is True + assert metadata["needs_session_id"] is False + assert metadata["action_space"] == "joint_position" diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index d274eadda..fa0148599 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -270,12 +270,14 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): def test_build_request_clones_stage_defaults_before_protocol_fields(): default_params = OmniDiffusionSamplingParams( + num_inference_steps=12, + output_type="latent", extra_args={ "format_prompt_as_json": True, "session_id": "configured-session-must-not-win", "reset": False, "nested": {"values": []}, - } + }, ) od_config = SimpleNamespace(model_config={"policy_server_config": {}}) engine_client = SimpleNamespace( @@ -297,10 +299,18 @@ def test_build_request_clones_stage_defaults_before_protocol_fields(): session_id="wire-session", reset=False, ) + request_c = serving._build_request( + {"prompt": "pick up the object", "sampling_params": {"num_inference_steps": 4}}, + session_id="wire-session", + reset=False, + ) assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True assert request_a.sampling_params.extra_args["session_id"] == "wire-session" assert request_a.sampling_params.extra_args["reset"] is True + assert request_a.sampling_params.num_inference_steps == 12 + assert request_a.sampling_params.output_type == "latent" + assert request_c.sampling_params.num_inference_steps == 4 request_a.sampling_params.extra_args["nested"]["values"].append("request-a") assert request_b.sampling_params.extra_args["nested"] == {"values": []} assert default_params.extra_args["nested"] == {"values": []} @@ -318,6 +328,46 @@ def test_build_request_forwards_seed_to_sampling_params(): assert isinstance(unseeded.sampling_params.seed, int) +# Per-request inference parameters live in their own namespace so robot feature +# names stay open-ended. Unknown keys inside it are ignored rather than rejected. +def test_build_request_forwards_num_inference_steps(): + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request( + {"prompt": "pick up the object", "sampling_params": {"num_inference_steps": 4}}, + session_id="session-a", + reset=True, + ) + + assert request.sampling_params.num_inference_steps == 4 + # The namespace is an engine concern; the pipeline must not see it as a feature. + assert "sampling_params" not in request.sampling_params.extra_args["robot_obs"] + + +def test_build_request_omits_num_inference_steps_when_not_sent(): + """Absent must stay absent: the pipeline reads None as "use the configured + default", so materialising a null here would be a different request.""" + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request({"prompt": "x"}, session_id="s", reset=True) + + assert request.sampling_params.num_inference_steps is None + + +def test_build_request_does_not_forward_arbitrary_observation_keys(): + """A robot observation may carry any feature name; none of them become + engine parameters.""" + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request( + {"prompt": "x", "state": np.zeros(32), "guidance_scale": 7.5}, session_id="s", reset=True + ) + + extra_args = request.sampling_params.extra_args + assert set(extra_args) == {"reset", "session_id", "robot_obs"} + assert extra_args["robot_obs"]["guidance_scale"] == 7.5 + + def test_infer_keeps_session_state_but_uses_unique_engine_request_ids(): engine = RecordingEngine() serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 0632db55d..975a6ddd2 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1166,6 +1166,7 @@ def pi0_openpi_run_policy_session( prompt: str = PI0_OPENPI_DEFAULT_PROMPT, session_id: str | None = None, num_steps: int = 2, + num_inference_steps: int | None = None, ) -> dict[str, Any]: """Connect, read handshake metadata, send ``num_steps`` observations.""" import uuid @@ -1184,6 +1185,8 @@ def pi0_openpi_run_policy_session( actions = [] for _ in range(num_steps): payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id) + if num_inference_steps is not None: + payload["sampling_params"] = {"num_inference_steps": num_inference_steps} payload["endpoint"] = "infer" conn.send(packer.pack(payload)) actions.append(_pi0_decode_action_response(conn.recv())) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f0a1285bc..ed266d734 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -82,6 +82,7 @@ EXCLUDED_MODELS = [ "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", "Pi0Pipeline", + "Pi05Pipeline", "SanaWmPipeline", # Audio pipeline conditioned by a separate Qwen2.5-Omni encoder stage; # covered by tests/diffusion/models/auk (CPU pipeline tests, parity) and diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 92aad9f4d..eada1019d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -33,6 +33,7 @@ from vllm_omni.config.stage_config import ( PipelineConfig, ) from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE +from vllm_omni.diffusion.models.pi05_pipeline_config import PI05_PIPELINE from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_S2S_PIPELINE, AUDEX_THINKER_ONLY_PIPELINE, @@ -153,6 +154,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, + "pi05": PI05_PIPELINE, # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 # checkpoints (which stay on the single-stage diffusion fallback), so this # entry is only reachable through a deploy yaml's ``pipeline:`` key diff --git a/vllm_omni/deploy/pi05.yaml b/vllm_omni/deploy/pi05.yaml new file mode 100644 index 000000000..aec8e88f0 --- /dev/null +++ b/vllm_omni/deploy/pi05.yaml @@ -0,0 +1,51 @@ +# π0.5 VLA: one diffusion stage, one GPU, eager, bsz 1 (OpenPI policy serving is +# one observation -> one action reply). Every key below is documented in +# recipes/lerobot/Pi05.md. + +pipeline: pi05 +async_chunk: false +distributed_executor_backend: mp +# float32 or bfloat16. bfloat16 is faster and halves memory at a numerical cost; +# read "Choosing a dtype" in the recipe before switching. +dtype: float32 + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + model_class_name: Pi05Pipeline + model_config: + paligemma_variant: gemma_2b + action_expert_variant: gemma_300m + chunk_size: 50 + # Must equal chunk_size: this endpoint returns the whole chunk and keeps no + # action queue, so "execute the first N steps, then replan" is not available. + n_action_steps: 50 + max_action_dim: 32 + max_state_dim: 32 + num_inference_steps: 10 + image_resolution: [224, 224] + tokenizer_max_length: 200 + state_num_bins: 256 + max_cameras: 3 + image_feature_keys: + - observation.images.base_0_rgb + - observation.images.left_wrist_0_rgb + - observation.images.right_wrist_0_rgb + image_key_map: {} + use_relative_actions: false + relative_exclude_joints: + - gripper + # OpenPI handshake metadata advertised to the robot client. + policy_server_config: + image_resolution: [224, 224] + n_external_cameras: 1 + needs_wrist_camera: true + needs_stereo_camera: false + needs_session_id: false + action_space: joint_position + action_horizon: 50 + action_dim: 32 + max_action_dim: 32 + max_cameras: 3 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0ffd0c54e..6aa1a27d9 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1615,6 +1615,11 @@ class OmniDiffusionConfig: self.model_class_name = "Pi0Pipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() + elif cfg.get("type") == "pi05": + if self.model_class_name is None: + self.model_class_name = "Pi05Pipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() elif architectures and len(architectures) == 1: architecture = architectures[0] from vllm_omni.diffusion.registry import DiffusionModelRegistry diff --git a/vllm_omni/diffusion/models/pi05/__init__.py b/vllm_omni/diffusion/models/pi05/__init__.py new file mode 100644 index 000000000..ccd19acde --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/__init__.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 VLA model for vllm-omni. + +PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert with AdaRMS +timestep conditioning + a flow-matching action head. Outputs a continuous action +chunk ``[horizon, action_dim]`` rather than tokens. + +Differs from π0 in five places: AdaRMS timestep conditioning (``time_mlp_*`` +instead of ``action_time_mlp_*``), AdaRMS norms throughout the action expert, a +200-token tokenizer budget, a discretized state carried in the prompt (no +``state_proj``), and relative-action support. +""" + +from vllm_omni.diffusion.models.pi05.config import ( + Pi05Config, + UnsupportedCheckpointCapabilityError, + resolve_excluded_action_indices, +) +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + GemmaVariantConfig, + PaliGemmaWithActionExpertPi05, + Pi05AdaRMSNorm, + Pi05ForActionPrediction, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, + prepare_attention_masks_4d, +) +from vllm_omni.diffusion.models.pi05.pipeline_pi05 import ( + Pi05Pipeline, + get_pi05_post_process_func, +) +from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05ImageProcessor, + Pi05Processor, + Pi05RelativeActions, + apply_norm, + build_norm_stats, + build_pi05_prompt, + discretize_state, + pil_image_to_tensor, + resize_with_pad, + tokenize_prompt, +) + +__all__ = [ + "Pi05Config", + "UnsupportedCheckpointCapabilityError", + "resolve_excluded_action_indices", + "Pi05ForActionPrediction", + "PaliGemmaWithActionExpertPi05", + "Pi05AdaRMSNorm", + "GemmaVariantConfig", + "get_gemma_config", + "create_sinusoidal_pos_embedding", + "make_att_2d_masks", + "prepare_attention_masks_4d", + "Pi05ImageProcessor", + "Pi05Processor", + "Pi05RelativeActions", + "build_pi05_prompt", + "discretize_state", + "apply_norm", + "build_norm_stats", + "pil_image_to_tensor", + "resize_with_pad", + "tokenize_prompt", + "Pi05Pipeline", + "get_pi05_post_process_func", +] diff --git a/vllm_omni/diffusion/models/pi05/config.py b/vllm_omni/diffusion/models/pi05/config.py new file mode 100644 index 000000000..e86ddbae9 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/config.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Config surface for the π0.5 VLA model in vllm-omni. + +Deliberately shaped exactly like ``pi0/config.py``: a small dataclass that +consumes the raw LeRobot ``config.json`` (the field surface of +``lerobot.policies.pi05.PI05Config``) and keeps only the runtime-relevant +fields. Transformer dimensions are derived from ``paligemma_variant`` / +``action_expert_variant`` inside the model via ``get_gemma_config``. + +What π0.5 adds on top of the π0 config surface: + +* ``tokenizer_max_length = 200`` (π0 uses 48). +* ``state_num_bins`` — state is discretized into language tokens instead of + going through a ``state_proj`` layer. +* ``use_relative_actions`` / ``relative_exclude_joints`` / + ``action_feature_names`` — the relative-action contract. See + ``processor_pi05.Pi05RelativeActions``. +* Quantile normalization stats. LeRobot's π0.5 defaults ``STATE`` and + ``ACTION`` to ``NormalizationMode.QUANTILES`` where π0 uses ``MEAN_STD``. + +**Checkpoint boundary rule.** A capability that the checkpoint *declares* but +this implementation does not *consume* must raise, not be silently dropped. +π0.5 checkpoints can declare MEM (short-horizon observation memory) and RTC +(real-time chunking); neither is supported here, and serving such a checkpoint +anyway would produce plausible-looking wrong actions. See +``_reject_unsupported_capabilities``. +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Callable +from dataclasses import dataclass, field +from dataclasses import fields as dataclass_fields +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +# Serving dtypes, by name. Kept in sync with ``Pi05Pipeline.SUPPORTED_DTYPES``, +# which guards the path that performs the actual cast. Defined here rather than +# imported to keep config.py free of a pipeline import cycle. +SUPPORTED_DTYPE_NAMES = frozenset({"float32", "bfloat16"}) + +# LeRobot / OpenPI observation key conventions. +ACTION = "action" +OBS_STR = "observation" +OBS_STATE = OBS_STR + ".state" +OBS_IMAGES = OBS_STR + ".images" + + +class UnsupportedCheckpointCapabilityError(ValueError): + """A checkpoint declares a capability this implementation does not consume. + + Raised at load time rather than silently ignored: every one of these + capabilities changes what a *correct* action chunk looks like, and none of + them is visible in the weights alone. + """ + + +def resolve_excluded_action_indices( + exclude_joints: list[str] | None, + action_names: list[str] | None, +) -> list[int]: + """Map ``relative_exclude_joints`` names onto action-vector indices. + + Matching is exact name first, then substring (a checkpoint may name the + gripper dimension ``gripper_position`` while the config just says + ``gripper``). Single source of truth for both the config-time validation + and the runtime mask in ``processor_pi05.Pi05RelativeActions``. + + Returns an empty list when there is nothing to exclude. Raises when a name + cannot be resolved — an unresolvable exclusion would otherwise silently + become "make this dimension relative too". + """ + if not [name for name in (exclude_joints or []) if name]: + return [] + if not action_names: + raise UnsupportedCheckpointCapabilityError( + f"Cannot resolve relative_exclude_joints={exclude_joints!r} without action_feature_names." + ) + + # LeRobot walks the action names and excludes each one any token equals or is + # contained in, both lowercased. Walking the tokens instead and stopping at the + # first exact hit misses a second name the token is a substring of. + tokens = [str(name).lower() for name in exclude_joints if name] + lowered = [str(name).lower() for name in action_names] + indices = [i for i, name in enumerate(lowered) if any(t == name or t in name for t in tokens)] + unresolved = [t for t in tokens if not any(t == name or t in name for name in lowered)] + + if unresolved: + raise UnsupportedCheckpointCapabilityError( + f"relative_exclude_joints entries {unresolved!r} match no entry in action_feature_names={action_names!r}." + ) + return sorted(set(indices)) + + +@dataclass +class Pi05Config: + """π0.5 VLA config (dataclass, not an HF ``PretrainedConfig``).""" + + # Backbone variants — mapped to Gemma dimensions by ``get_gemma_config``. + paligemma_variant: str = "gemma_2b" + action_expert_variant: str = "gemma_300m" + + # Action chunk shape. + chunk_size: int = 50 + # The stateless OpenPI endpoint returns one complete predicted chunk. + n_action_steps: int = 50 + max_action_dim: int = 32 + max_state_dim: int = 32 + + # Flow-matching denoising schedule. + num_inference_steps: int = 10 + # Sinusoidal timestep embedding periods (must match OpenPI/LeRobot). + min_period: float = 4e-3 + max_period: float = 4.0 + + # Image preprocessing. π0.5/SigLIP only support square inputs. + image_resolution: tuple[int, int] = (224, 224) + # π0.5 pads text to 200 tokens (π0 uses 48). The prompt now also carries the + # serialized state, which is why it is so much longer. + tokenizer_max_length: int = 200 + # Number of camera slots the model attends to (real + padded). + max_cameras: int = 3 + + # Bins the normalized state is discretized into before it enters the prompt. + # Ref: openpi PaliGemmaTokenizer.tokenize(). + state_num_bins: int = 256 + + # Weight dtype the checkpoint was saved in. + dtype: str = "float32" + + # ── Relative actions ────────────────────────────────────────────── + # True when the checkpoint was trained on actions relative to the current + # state. Its ``norm_stats`` are then in relative space, so serving it without + # the transform silently yields wrong actions — the weights look identical + # either way. + use_relative_actions: bool = False + # Joint names kept absolute (gripper open/close is an absolute quantity). + relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"]) + # Needed to resolve ``relative_exclude_joints``; there is no dataset to + # fall back on at serving time. + action_feature_names: list[str] | None = None + + # Per-dataset normalization stats (schema matches LeRobot's + # ``NormalizerProcessorStep``). ``None`` means identity / pass-through. + # π0.5 defaults to quantile mode; see ``processor_pi05.build_norm_stats``. + norm_stats: dict | None = None + + # Ordered camera-slot identities. Missing keys retain their slot and are + # represented by an empty image with a false mask during preprocessing. + image_feature_keys: list[str] | None = None + # Optional map from raw OpenPI obs keys → ``image_feature_keys`` entries. + image_key_map: dict[str, str] = field(default_factory=dict) + + # Checkpoint feature schemas. The input schema determines camera order and + # the state width serialized into the prompt; the output schema determines + # the unpadded action width returned on the wire. + input_features: dict[str, Any] = field(default_factory=dict) + output_features: dict[str, Any] = field(default_factory=dict) + # OpenPI handshake metadata from the deploy config. Construction validates + # it against the resolved checkpoint contract. + policy_server_config: dict[str, Any] = field(default_factory=dict) + + # Derived from output_features[ACTION].shape; never accepted as a second + # source of truth in config.json. + action_dim: int = field(init=False) + # Derived from input_features[OBS_STATE].shape — the number of values the + # prompt serializes, not a padding target. See processor_pi05.as_state_vector. + state_dim: int = field(init=False) + + def __post_init__(self) -> None: + # Coerce list → tuple (JSON has no tuples). Squareness is a real + # constraint: SigLIP only takes square inputs. + res = self.image_resolution + if not isinstance(res, (tuple, list)) or len(res) != 2 or res[0] != res[1]: + raise ValueError(f"π0.5 expects a square image_resolution (H == W); got {res!r}.") + self.image_resolution = (int(res[0]), int(res[1])) + + if self.n_action_steps != self.chunk_size: + raise UnsupportedCheckpointCapabilityError( + "The stateless OpenPI serving path returns one complete action chunk and " + f"does not support n_action_steps={self.n_action_steps} with " + f"chunk_size={self.chunk_size}." + ) + # Mirrors Pi05Pipeline.SUPPORTED_DTYPES. This field records what the + # *checkpoint* declares; the dtype the weights are actually cast to comes + # from the top-level OmniDiffusionConfig and is validated there. + if self.dtype not in SUPPORTED_DTYPE_NAMES: + raise ValueError(f"dtype must be one of {sorted(SUPPORTED_DTYPE_NAMES)}, got {self.dtype!r}.") + + # Derive the camera order from input_features if not given explicitly. + if self.image_feature_keys is None and self.input_features: + self.image_feature_keys = [key for key in self.input_features if key.startswith(OBS_IMAGES + ".")] + + if self.max_cameras < 1: + raise ValueError(f"max_cameras must be at least 1, got {self.max_cameras!r}.") + if self.num_inference_steps < 1: + # ``dt = -1 / num_steps`` in the Euler loop. + raise ValueError(f"num_inference_steps must be at least 1, got {self.num_inference_steps!r}.") + + self._derive_widths() + self._validate_norm_stats() + self._validate_policy_server_config() + self._validate_relative_actions() + + def _validate_norm_stats(self) -> None: + """Statistics have to cover the widths the checkpoint declares.""" + for key, width in (("state", self.state_dim), ("action", self.action_dim)): + for name, value in (self.norm_stats or {}).get(key, {}).items(): + if name != "mode" and len(value) != width: + raise ValueError(f"norm_stats[{key}][{name}] has {len(value)} entries, expected {width}.") + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + def _derive_widths(self) -> None: + """Read the real state and action widths off the checkpoint schema. + + π0.5 serializes the state into the prompt at its real width, so a 7-joint + checkpoint must not be padded out to ``max_state_dim``. The action chunk + is cropped to ``action_dim`` on the way out. + """ + state = self.input_features.get(OBS_STATE) or {} + action = self.output_features.get(ACTION) or {} + if state.get("type", "STATE").upper() != "STATE": + raise ValueError(f"input_features[{OBS_STATE!r}].type must be STATE, got {state['type']!r}.") + if action.get("type", "ACTION").upper() != "ACTION": + raise ValueError(f"output_features[{ACTION!r}].type must be ACTION, got {action['type']!r}.") + self.state_dim = int(state["shape"][0]) if state else self.max_state_dim + self.action_dim = int(action["shape"][0]) if action else self.max_action_dim + if self.state_dim > self.max_state_dim or self.action_dim > self.max_action_dim: + raise ValueError( + f"Checkpoint declares state_dim={self.state_dim} / action_dim={self.action_dim}, " + f"over max_state_dim={self.max_state_dim} / max_action_dim={self.max_action_dim}." + ) + + def _validate_policy_server_config(self) -> None: + """The OpenPI handshake has to advertise the resolved contract.""" + expected = { + "action_horizon": self.chunk_size, + "action_dim": self.action_dim, + "max_action_dim": self.max_action_dim, + "max_cameras": self.max_cameras, + "image_resolution": list(self.image_resolution), + } + for key, value in expected.items(): + declared = (self.policy_server_config or {}).get(key) + if declared is None: + continue + if key == "image_resolution": + declared = list(declared) + if declared != value: + raise ValueError( + f"policy_server_config.{key}={declared!r} does not match the resolved value {value!r}." + ) + + def _validate_relative_actions(self) -> None: + """Resolve the excluded joints at load, so a bad name fails before serving.""" + if not self.use_relative_actions or not self.relative_exclude_joints: + return + if not self.action_feature_names: + raise UnsupportedCheckpointCapabilityError( + f"relative_exclude_joints={self.relative_exclude_joints!r} needs action_feature_names, " + "which the checkpoint does not carry." + ) + resolve_excluded_action_indices(self.relative_exclude_joints, self.action_feature_names) + + # ------------------------------------------------------------------ + # Constructors + # ------------------------------------------------------------------ + @classmethod + def from_pretrained(cls, checkpoint_dir: str | Path) -> Pi05Config: + """Build from a checkpoint directory's ``config.json``. + + Normalization stats are *not* in ``config.json``. LeRobot keeps them in + the processor sidecar, so they are loaded separately and backfilled + here — see :func:`load_lerobot_norm_stats`. + """ + checkpoint_dir = Path(checkpoint_dir) + config_path = checkpoint_dir / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"π0.5 checkpoint is missing required config: {config_path}.") + with open(config_path, encoding="utf-8") as f: + raw = json.load(f) + config = cls.from_model_config(raw) + + if config.norm_stats is None: + stats = load_lerobot_norm_stats(checkpoint_dir) + if stats: + config.norm_stats = stats + config._validate_norm_stats() + return config + + @classmethod + def from_model_config(cls, model_config: dict[str, Any] | None) -> Pi05Config: + """Build from a config dict (LeRobot ``config.json`` or deploy yaml).""" + if not model_config: + return cls() + + raw = dict(model_config) + _reject_unsupported_capabilities(raw) + + model_type = raw.pop("type", "pi05") + if model_type != "pi05": + raise ValueError(f"Expected a π0.5 checkpoint (type='pi05'), got type={model_type!r}.") + + if "image_resolution" in raw: + raw["image_resolution"] = tuple(raw["image_resolution"]) + + # A config.json carries the whole training recipe; none of it reaches + # inference. Capabilities we cannot serve are rejected above instead. + allowed = {item.name for item in dataclass_fields(cls) if item.init} + filtered = {key: value for key, value in raw.items() if key in allowed} + dropped = sorted(set(raw) - allowed) + if dropped: + logger.debug("π0.5 config: ignoring %d non-runtime key(s): %s", len(dropped), dropped) + return cls(**filtered) + + +# ---------------------------------------------------------------------- +# LeRobot normalization-stats sidecar +# ---------------------------------------------------------------------- +# Stats live beside config.json: ``policy_preprocessor.json`` names a +# safetensors file per stateful step, keyed ``"."``. +_PREPROCESSOR_JSON = "policy_preprocessor.json" + +# The mode must come from ``norm_map``: a state_dict carries all of +# mean/std/min/max/q01/q99 regardless, so guessing from the present keys would +# read a QUANTILES checkpoint as mean_std and apply a wrong affine map. +_NORM_MODE_FROM_LEROBOT: dict[str, tuple[str, str, str] | None] = { + "IDENTITY": None, + "MEAN_STD": ("mean_std", "mean", "std"), + "MIN_MAX": ("min_max", "min", "max"), + "QUANTILES": ("quantile", "q01", "q99"), +} +# Renaming happens in an earlier processor step, so these names are canonical +# here. VISUAL is absent by design: images are normalized in the processor. +_NORM_STATS_FEATURES = {OBS_STATE: ("STATE", "state"), ACTION: ("ACTION", "action")} + + +def load_lerobot_norm_stats(checkpoint_dir: str | Path) -> dict[str, dict[str, Any]] | None: + """Load normalization stats from a LeRobot checkpoint's processor sidecar. + + Returns a ``norm_stats``-shaped dict (``{"state": {...}, "action": {...}}``) + carrying an **explicit** ``mode``, or ``None`` when the checkpoint ships no + stats. ``lerobot/pi05_base`` is the latter case: its normalizer step has no + ``state_file``, i.e. normalization is identity and the client is expected to + send an already-normalized state. + + Raises on a mode we cannot reproduce rather than serving the checkpoint with + the wrong transform, which fails silently — a wrongly normalized state still + yields a plausible-looking action chunk. + """ + checkpoint_dir = Path(checkpoint_dir) + preprocessor_path = checkpoint_dir / _PREPROCESSOR_JSON + if not preprocessor_path.exists(): + return None + with open(preprocessor_path, encoding="utf-8") as f: + steps = json.load(f).get("steps", []) + + step = next((s for s in steps if s.get("registry_name") == "normalizer_processor"), {}) + state_file = step.get("state_file") + if not state_file: + logger.info( + "π0.5 config: %s declares no normalizer state file; the state passes through.", + _PREPROCESSOR_JSON, + ) + return None + state_path = checkpoint_dir / state_file + if not state_path.exists(): + raise FileNotFoundError(f"{preprocessor_path} declares normalizer state {state_file!r}, which is missing.") + + norm_map = {str(k).upper(): str(v).upper() for k, v in ((step.get("config") or {}).get("norm_map") or {}).items()} + + import safetensors.torch + + flat = safetensors.torch.load_file(str(state_path)) + + stats: dict[str, dict[str, Any]] = {} + for feature_name, (feature_type, stats_key) in _NORM_STATS_FEATURES.items(): + lerobot_mode = norm_map.get(feature_type, "IDENTITY") + if lerobot_mode not in _NORM_MODE_FROM_LEROBOT: + raise ValueError( + f"π0.5 cannot reproduce normalization mode {lerobot_mode!r} for {feature_type}. " + f"Expected one of {sorted(_NORM_MODE_FROM_LEROBOT)}." + ) + selected = _NORM_MODE_FROM_LEROBOT[lerobot_mode] + if selected is None: + continue + + mode, *stat_names = selected + missing = [name for name in stat_names if f"{feature_name}.{name}" not in flat] + if missing: + raise ValueError( + f"{feature_type} declares {lerobot_mode} normalization but the state file is missing {missing}." + ) + stats[stats_key] = {"mode": mode} | {name: flat[f"{feature_name}.{name}"].tolist() for name in stat_names} + + if not stats: + return None + logger.info( + "π0.5 config: loaded normalization stats from %s — %s.", + state_file, + ", ".join(f"{key}={value['mode']}" for key, value in sorted(stats.items())), + ) + return stats + + +# ---------------------------------------------------------------------- +# Checkpoint-boundary rule +# ---------------------------------------------------------------------- +# Each entry: config key → (is this value *enabled*?, why we cannot serve it). +# Every capability here changes what a *correct* action chunk looks like, and +# none is visible in the weights, so a checkpoint that enables one is refused +# rather than served plausibly wrong. See recipes/lerobot/Pi05.md for the +# longer explanations. +# +# The predicate is per key because these are not all flags: ``n_obs_steps`` is +# a count whose supported value is 1, and an empty ``rtc_config`` still selects +# RTC — LeRobot's ``RTCConfig`` defaults ``enabled=True``. +_UNSUPPORTED: dict[str, tuple[Callable[[Any], bool], str]] = { + "use_peft": (bool, "adapter-aware loading; this loader takes a merged checkpoint only"), + "empty_cameras": (bool, "mutates the input schema; declare camera features explicitly instead"), + "use_visual_memory": (bool, "MEM needs a per-session observation history, which this path does not keep"), + "use_proprioceptive_memory": ( + bool, + "MEM sends the state as a projected token, not the discretized prompt built here", + ), + "rtc_config": ( + lambda value: value is not None, + "RTC needs prefix guidance in the denoising loop and per-request chunk carry-over", + ), + "n_obs_steps": ( + lambda value: bool(value) and value > 1, + "an observation history; this path is first-order Markov", + ), +} + + +def _reject_unsupported_capabilities(raw: dict[str, Any]) -> None: + """Raise if the checkpoint enables something we do not consume. + + Scope note: this validates that we correctly consume *the checkpoint we were + handed*. Choosing which checkpoint to load belongs to the caller. + """ + problems = [ + f" - {key}={raw[key]!r}: {why}" + for key, (is_enabled, why) in _UNSUPPORTED.items() + if key in raw and is_enabled(raw[key]) + ] + if problems: + raise UnsupportedCheckpointCapabilityError( + "This checkpoint declares capabilities the vllm-omni π0.5 implementation " + "does not support:\n" + "\n".join(problems) + "\nServing it anyway would " + "produce plausible-looking but wrong actions." + ) + + # Not an error: it records that the checkpoint was trained with clean action + # prefixes sampled in, which is a property of training rather than a request + # to run RTC. Such a checkpoint is still correct to serve without RTC. + delay = raw.get("rtc_training_max_delay") + if delay: + logger.info( + "π0.5 config: checkpoint was trained with rtc_training_max_delay=%s. " + "Real-Time Chunking is not implemented here; serving proceeds without it.", + delay, + ) diff --git a/vllm_omni/diffusion/models/pi05/modeling_pi05.py b/vllm_omni/diffusion/models/pi05/modeling_pi05.py new file mode 100644 index 000000000..a792ae0ef --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/modeling_pi05.py @@ -0,0 +1,898 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Inference-only π0.5 VLA math kernel for vllm-omni. + +Only the math that turns a robot observation into an action chunk; no serving or +request glue. Deliberately shaped like ``models/pi0/modeling_pi0.py`` so the two +can later be factored into a shared Pi-family module (RFC step 2) on a +"behaviour unchanged" review. + +π0.5 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert +suffix + flow-matching head. Where π0 projects the robot state through +``state_proj``, π0.5 discretizes it into prompt tokens — so there is no +``state_proj`` layer, ``sample_actions`` takes no ``state`` argument, and the +suffix is action tokens only, which drops π0's leading state-token boundary from +the suffix attention mask and leaves ``[1] + [0] * (horizon - 1)``. + +π0.5 is nonetheless the *larger* model: the 37 AdaRMS ``dense`` projections add +~116M parameters against the ~8K that ``state_proj`` saves. (LeRobot's README +says otherwise; the checkpoint disagrees.) + +Reference implementations: + - OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py + - LeRobot: lerobot/src/lerobot/policies/pi05/modeling_pi05.py +""" + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.gemma.modeling_gemma import ( + GemmaForCausalLM, + apply_rotary_pos_emb, +) +from transformers.models.paligemma.modeling_paligemma import ( + PaliGemmaForConditionalGeneration, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +logger = logging.getLogger(__name__) + +# ────────────────────────────────────────────────────────────────────── +# Constants +# ────────────────────────────────────────────────────────────────────── +DEFAULT_ACTION_DIM = 32 +DEFAULT_ACTION_HORIZON = 50 +DEFAULT_MAX_TOKEN_LEN = 200 # π0 uses 48 +DEFAULT_NUM_INFERENCE_STEPS = 10 +DEFAULT_IMAGE_RESOLUTION = (224, 224) # openpi/models/model.py IMAGE_RESOLUTION +DEFAULT_STATE_NUM_BINS = 256 # openpi PaliGemmaTokenizer.tokenize() + +# Large negative value to fill masked-out positions in a float attention mask. +# Matches OpenPI's constant exactly so that numerics line up during parity. +# Ref: openpi/src/openpi/models/gemma.py +OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38 + + +# ────────────────────────────────────────────────────────────────────── +# Gemma variant configs (matches openpi/models/gemma.py get_config) +# ────────────────────────────────────────────────────────────────────── +class GemmaVariantConfig: + def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim): + self.width = width + self.depth = depth + self.mlp_dim = mlp_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + + +def get_gemma_config(variant: str) -> GemmaVariantConfig: + if variant == "gemma_2b": + return GemmaVariantConfig(2048, 18, 16384, 8, 1, 256) + elif variant == "gemma_300m": + return GemmaVariantConfig(1024, 18, 4096, 8, 1, 256) + else: + raise ValueError(f"Unknown variant: {variant}") + + +# ────────────────────────────────────────────────────────────────────── +# Utility functions (match openpi/models_pytorch/pi0_pytorch.py) +# ────────────────────────────────────────────────────────────────────── +def create_sinusoidal_pos_embedding( + time: torch.Tensor, + dimension: int, + min_period: float = 4e-3, + max_period: float = 4.0, + device: torch.device = None, +) -> torch.Tensor: + """Compute a sine/cosine positional embedding for scalar timesteps. + + Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding + """ + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + if time.ndim != 1: + raise ValueError("time tensor must be 1-D (batch_size,)") + if device is None: + device = time.device + + # Use float64 for the log-linear sweep and for the inner products, to + # match the reference implementation's numerical behaviour exactly. + dtype = torch.float64 + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) + period = min_period * (max_period / min_period) ** fraction + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None].to(dtype) + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: + """Build a 2D attention mask from a padding mask and an autoregressive mask. + + Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks + """ + if att_masks.ndim != 2: + raise ValueError(f"att_masks must be 2-D, got {att_masks.ndim}-D") + if pad_masks.ndim != 2: + raise ValueError(f"pad_masks must be 2-D, got {pad_masks.ndim}-D") + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d_masks & pad_2d_masks + + +def prepare_attention_masks_4d(att_2d_masks: torch.Tensor) -> torch.Tensor: + """Convert ``(B, S, S)`` bool masks to ``(B, 1, S, S)`` float masks. + + ``True`` → 0.0 (attend), ``False`` → ``OPENPI_ATTENTION_MASK_VALUE``. + """ + att_2d_masks_4d = att_2d_masks[:, None, :, :] + return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE) + + +class Pi05AdaRMSNorm(nn.Module): + """Adaptive RMSNorm conditioned on the flow-matching timestep. + + π0 conditions on time by *concatenating* a time embedding onto each action + embedding. π0.5 instead feeds the time embedding into every action-expert + norm, which produces a per-layer ``(scale, shift, gate)`` triple:: + + y = norm(x) * (1 + scale) + shift + out = residual + gate * sublayer(y) + + ``dense`` is zero-initialized, so an untrained model starts as the identity + modulation with a closed gate — matching OpenPI's parameterization. + + Note the shape of the unconditioned branch: ``normed * (1 + weight)`` with + ``weight`` zero-initialized, which is exactly ``transformers``' + ``GemmaRMSNorm``. That equivalence is why only the *expert* norms need + replacing here and the PaliGemma prefix can keep stock Gemma layers. + """ + + def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None): + super().__init__() + self.eps = eps + self.dim = dim + self.cond_dim = cond_dim + if cond_dim is not None: + self.dense = nn.Linear(cond_dim, dim * 3, bias=True) + nn.init.zeros_(self.dense.weight) + nn.init.zeros_(self.dense.bias) + self.weight = None + else: + self.weight = nn.Parameter(torch.zeros(dim)) + self.dense = None + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + return x.float() * torch.rsqrt(var + self.eps) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return ``(normed, gate)``; ``gate`` is ``None`` in the unconditioned case.""" + dtype = x.dtype + normed = self._norm(x) + if self.dense is None: + normed = normed * (1.0 + self.weight.float()) + return normed.to(dtype), None + + if cond is None: + # A conditioned norm silently falling back to an unconditioned one + # would drop the entire timestep signal and still return a + # well-shaped, finite tensor. + raise ValueError( + "Pi05AdaRMSNorm was built with cond_dim=" + f"{self.cond_dim} but called without an AdaRMS conditioning vector." + ) + + if cond.shape[-1] != self.cond_dim: + raise ValueError(f"Expected AdaRMS cond dim {self.cond_dim}, got {cond.shape[-1]}") + + modulation = self.dense(cond.to(self.dense.weight.dtype)) + if x.ndim == 3: + # (B, 3*dim) → (B, 1, 3*dim), broadcast across the token axis: the + # timestep is a per-sample scalar, identical for every action token. + modulation = modulation.unsqueeze(1) + scale, shift, gate = modulation.chunk(3, dim=-1) + normed = normed * (1.0 + scale.float()) + shift.float() + return normed.to(dtype), gate.to(dtype) + + +def _gated_residual(residual: torch.Tensor, out: torch.Tensor, gate: torch.Tensor | None) -> torch.Tensor: + if gate is None: + return residual + out + return residual + gate * out + + +# ────────────────────────────────────────────────────────────────────── +# Dual-backbone: PaliGemma + AdaRMS action expert +# ────────────────────────────────────────────────────────────────────── +def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention. + + ``(B, num_kv_heads, S, D) → (B, num_kv_heads * n_rep, S, D)``. + """ + if n_rep == 1: + return hidden_states + b, nh, s, d = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(b, nh, n_rep, s, d) + return hidden_states.reshape(b, nh * n_rep, s, d) + + +def _attend(query_states, key_states, value_states, attention_mask, num_kv_groups, scaling): + """Manual eager attention: ``softmax(Q Kᵀ · scale + mask) · V``. + + The mask is sliced to ``key_states.shape[-2]`` so the same + ``(B, 1, Q, prefix+suffix)`` mask works in both the prefix pass (K length = + prefix) and the suffix pass (K length = prefix + suffix). + """ + k = _repeat_kv(key_states, num_kv_groups) + v = _repeat_kv(value_states, num_kv_groups) + attn_weights = torch.matmul(query_states, k.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask[:, :, :, : k.shape[-2]] + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + return torch.matmul(attn_weights, v) + + +def _match(tensor: torch.Tensor, module: nn.Module) -> torch.Tensor: + """Cast ``tensor`` to the dtype ``module``'s weight expects.""" + return tensor.to(module.weight.dtype) if tensor.dtype != module.weight.dtype else tensor + + +def _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma): + """Run one PaliGemma LM layer on the prefix, returning the layer output and + the post-RoPE ``(k, v)`` for the suffix pass. + + Identical to π0: the prefix backbone is unchanged in π0.5 (no AdaRMS — + there is no timestep in the prefix). + """ + model = paligemma.model.language_model + layer = model.layers[layer_idx] + residual = hidden_states + x = _match(layer.input_layernorm(hidden_states), layer.self_attn.q_proj) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + cos, sin = model.rotary_emb(v, position_ids) + q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + out = layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)) + residual + after_resid = out + normed = layer.post_attention_layernorm(out) + out = layer.mlp(_match(normed, layer.mlp.up_proj)) + after_resid + return out, (k, v) + + +def _compute_layer_suffix_only( + layer_idx, + hidden_states, + prefix_kv, + attention_mask, + position_ids, + gemma_expert, + adarms_cond, +): + """Run one action-expert layer on the suffix with AdaRMS conditioning. + + This is where π0.5 diverges from π0. Both norms are + :class:`Pi05AdaRMSNorm` and each returns a gate that scales its sublayer's + contribution to the residual stream. + """ + layer = gemma_expert.model.layers[layer_idx] + + residual = hidden_states + x, gate = layer.input_layernorm(hidden_states, adarms_cond) + x = _match(x, layer.self_attn.q_proj) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k_suf = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v_suf = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + # RoPE frequencies are shared between PaliGemma and the expert. + cos, sin = gemma_expert.model.rotary_emb(v_suf, position_ids) + q, k_suf = apply_rotary_pos_emb(q, k_suf, cos, sin, unsqueeze_dim=1) + + # Concatenate cached prefix K/V (possibly different dtype) with suffix K/V. + k_prefix, v_prefix = prefix_kv + k = torch.cat([k_prefix.to(k_suf.dtype), k_suf], dim=2) + v = torch.cat([v_prefix.to(v_suf.dtype), v_suf], dim=2) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + hidden_states = _gated_residual(residual, layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)), gate) + + residual = hidden_states + x, gate = layer.post_attention_layernorm(hidden_states, adarms_cond) + return _gated_residual(residual, layer.mlp(_match(x, layer.mlp.up_proj)), gate) + + +class PaliGemmaWithActionExpertPi05(nn.Module): + """Dual-backbone transformer: PaliGemma (Gemma 2B) + AdaRMS expert (300M). + + Same two-mode dispatch as π0 (``prefix_only`` / ``suffix_only``), with one + structural change: after building a stock ``GemmaForCausalLM`` expert, every + norm in it is swapped for a :class:`Pi05AdaRMSNorm` carrying a ``dense`` + conditioning projection. + + Swapping in place — rather than subclassing ``GemmaModel`` as #4419 does — + keeps the module tree, and therefore the checkpoint key layout, identical to + the expert's stock layout apart from the norms themselves. + """ + + def __init__(self, vlm_config, action_expert_config): + super().__init__() + + # PaliGemma prefix: identical to π0. It sees no timestep, so no AdaRMS. + vlm_config_hf = CONFIG_MAPPING["paligemma"]() + vlm_config_hf._vocab_size = 257152 + vlm_config_hf.image_token_index = 257152 + vlm_config_hf.text_config.hidden_size = vlm_config.width + vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads + vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh" + vlm_config_hf.text_config.dtype = "float32" + vlm_config_hf.text_config.vocab_size = 257152 + vlm_config_hf.vision_config.intermediate_size = 4304 + vlm_config_hf.vision_config.projection_dim = 2048 + vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast" + vlm_config_hf.vision_config.dtype = "float32" + + action_expert_config_hf = CONFIG_MAPPING["gemma"]( + head_dim=action_expert_config.head_dim, + hidden_size=action_expert_config.width, + intermediate_size=action_expert_config.mlp_dim, + num_attention_heads=action_expert_config.num_heads, + num_hidden_layers=action_expert_config.depth, + num_key_value_heads=action_expert_config.num_kv_heads, + vocab_size=257152, + hidden_activation="gelu_pytorch_tanh", + dtype="float32", + ) + + self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf) + self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) + # The action expert doesn't embed tokens — it only consumes the + # suffix action embeddings we feed in. + self.gemma_expert.model.embed_tokens = None + + self.adarms_cond_dim = action_expert_config.width + self._install_adarms_norms(action_expert_config) + + def _install_adarms_norms(self, action_expert_config) -> None: + """Replace every action-expert RMSNorm with a conditioned AdaRMS norm.""" + expert = self.gemma_expert.model + eps = getattr(self.gemma_expert.config, "rms_norm_eps", 1e-6) + width = action_expert_config.width + for layer in expert.layers: + layer.input_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + layer.post_attention_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + expert.norm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + + def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Encode images with SigLIP vision tower + PaliGemma projector. + + The two steps are run explicitly rather than via + ``PaliGemmaModel.get_image_features`` because that helper divides the + projector output by ``sqrt(text_hidden_size)``. Being explicit keeps the + scale unambiguous and matches π0 exactly (SigLIP is unchanged in π0.5). + """ + # Shapes: pixel_values (B, 3, 224, 224) → SigLIP (B, 256, 1152) + # → projector (B, 256, 2048) + vision_outputs = self.paligemma.model.vision_tower(pixel_values) + image_features = vision_outputs.last_hidden_state + return self.paligemma.model.multi_modal_projector(image_features) + + def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + """Embed language tokens, returning the ``* sqrt(hidden)``-scaled embedding. + + The scaling location moved across transformers releases: at ≤5.3 the + normalizer lives inside ``GemmaModel.forward`` (which we bypass), and at + ≥5.4 ``GemmaTextScaledWordEmbedding`` self-applies it. Detect and avoid + double-scaling. + """ + embed_tokens = self.paligemma.model.language_model.embed_tokens + lang_emb = embed_tokens(tokens) + if getattr(embed_tokens, "embed_scale", None) is None: + lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1]) + return lang_emb + + def forward( + self, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values=None, + inputs_embeds: list[torch.Tensor | None] | None = None, + use_cache: bool = False, + adarms_cond: torch.Tensor | None = None, + ): + """Dispatch to prefix_only / suffix_only and return + ``([prefix_out, suffix_out], past_key_values_or_None)``. + """ + num_layers = self.paligemma.config.text_config.num_hidden_layers + pali_lm = self.paligemma.model.language_model + expert_lm = self.gemma_expert.model + + if inputs_embeds[1] is None: + hidden_states = inputs_embeds[0] + kv_list: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer_idx in range(num_layers): + hidden_states, kv = _compute_layer_prefix_only( + layer_idx, + hidden_states, + attention_mask, + position_ids, + paligemma=self.paligemma, + ) + kv_list.append(kv) + hidden_states = pali_lm.norm(hidden_states) + return [hidden_states, None], (kv_list if use_cache else None) + + if inputs_embeds[0] is not None: + raise ValueError( + "PaliGemmaWithActionExpertPi05.forward only supports prefix-only " + "or suffix-only dispatch; got both inputs_embeds populated." + ) + if not isinstance(past_key_values, list): + raise TypeError( + "suffix_only forward expects past_key_values to be the " + "list[(k, v)] produced by a previous prefix_only forward; " + f"got {type(past_key_values)}" + ) + hidden_states = inputs_embeds[1] + for layer_idx in range(num_layers): + hidden_states = _compute_layer_suffix_only( + layer_idx, + hidden_states, + past_key_values[layer_idx], + attention_mask, + position_ids, + gemma_expert=self.gemma_expert, + adarms_cond=adarms_cond, + ) + hidden_states, _ = expert_lm.norm(hidden_states, adarms_cond) + return [None, hidden_states], None + + +# ────────────────────────────────────────────────────────────────────── +# Main π0.5 Model +# ────────────────────────────────────────────────────────────────────── +class Pi05ForActionPrediction(nn.Module): + """π0.5 VLA model for robot action prediction via flow matching. + + Inference flow: + 1. Embed prefix (images + language, where the language already carries the + discretized state) → prefix tokens. + 2. Forward prefix through PaliGemma → layer-wise KV cache. + 3. For each denoising step ``t = 1.0, 1-dt, ..., 0``: + a. Embed the timestep → an AdaRMS conditioning vector. + b. Embed the suffix (action tokens only). + c. Forward the suffix through the AdaRMS action expert. + d. ``x_t = x_t + dt * v_t`` (Euler integration). + 4. Return ``x_0`` as the predicted action chunk. + """ + + def __init__(self, config, quant_config=None, prefix: str = ""): + super().__init__() + # ``quant_config`` is accepted for interface compatibility but unused — + # quant_config is not plumbed through; the weight dtype comes from the pipeline. + del quant_config + self.config = config + + self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM) + self.max_state_dim = getattr(config, "max_state_dim", self.action_dim) + self.action_horizon = getattr(config, "chunk_size", DEFAULT_ACTION_HORIZON) + self.num_inference_steps = getattr(config, "num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS) + + paligemma_variant = getattr(config, "paligemma_variant", "gemma_2b") + action_expert_variant = getattr(config, "action_expert_variant", "gemma_300m") + vlm_config = get_gemma_config(paligemma_variant) + expert_config = get_gemma_config(action_expert_variant) + self.vlm_width = vlm_config.width + self.expert_width = expert_config.width + + # Dual backbone + self.paligemma_with_expert = PaliGemmaWithActionExpertPi05(vlm_config, expert_config) + + # Action chunk projections. + self.action_in_proj = nn.Linear(self.action_dim, self.expert_width) + self.action_out_proj = nn.Linear(self.expert_width, self.action_dim) + + # π0.5 timestep MLP: (W → W → W) with SiLU, feeding AdaRMS. + # π0 instead has action_time_mlp_{in,out} of shape (2W → W → W) because + # it concatenates the time embedding onto the action embedding. + # NOTE: there is deliberately **no** ``state_proj`` here — that is the + # π0-only continuous-state path. + self.time_mlp_in = nn.Linear(self.expert_width, self.expert_width) + self.time_mlp_out = nn.Linear(self.expert_width, self.expert_width) + + # ── Prefix embedding ───────────────────────────────────────────── + def embed_prefix( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build prefix embeddings, per-token padding mask, and AR mask. + + Prefix tokens form ``[img_cam_0..., ..., lang_tokens...]`` with fully + bidirectional attention (all-zero ``att_masks``). Identical to π0 — the + state is inside ``lang_tokens``, so nothing here changes shape-wise. + + Cameras are embedded one at a time; each call is ``(B, 3, 224, 224)``. + The number of slots is fixed by ``config.max_cameras`` for the deployed + model; missing cameras occupy their slot with a false image mask. + """ + num_views = len(images) + if len(image_masks) != num_views: + raise ValueError( + f"images and image_masks must contain the same number of views, got {num_views} and {len(image_masks)}." + ) + max_cameras = int(self.config.max_cameras) + if num_views != max_cameras: + raise ValueError(f"Expected exactly max_cameras={max_cameras} image views, got {num_views}.") + + embs: list[torch.Tensor] = [] + pad_masks: list[torch.Tensor] = [] + att_masks: list[int] = [] + + for img, img_mask in zip(images, image_masks): + img_emb = self.paligemma_with_expert.embed_image(img) + bsize, num_img_embs = img_emb.shape[:2] + embs.append(img_emb) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + att_masks += [0] * num_img_embs + + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + embs.append(lang_emb) + pad_masks.append(lang_masks) + att_masks += [0] * lang_emb.shape[1] + + embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad_masks, dim=1) + att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device) + att_masks = att_masks[None, :].expand(pad_masks.shape[0], -1) + + return embs, pad_masks, att_masks + + # ── Timestep + suffix embedding ────────────────────────────────── + def embed_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + """Timestep → AdaRMS conditioning vector ``(B, expert_width)``. + + ``silu(time_mlp_out(silu(time_mlp_in(sinusoid(t)))))``. The trailing + SiLU is part of the reference implementation — dropping it is a silent + numerical error, not a crash. + """ + model_dtype = self.action_in_proj.weight.dtype + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=getattr(self.config, "min_period", 4e-3), + max_period=getattr(self.config, "max_period", 4.0), + device=timestep.device, + ).to(dtype=model_dtype) + time_cond = self.time_mlp_in(time_emb) + time_cond = F.silu(time_cond) + time_cond = self.time_mlp_out(time_cond) + return F.silu(time_cond) + + def embed_suffix( + self, + noisy_actions: torch.Tensor, + timestep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the suffix: **action tokens only**, plus the AdaRMS condition. + + π0's suffix is ``[state_token, action_tokens×H]`` with an AR mask of + ``[1, 1, 0...]``. π0.5 has no state token, so the suffix is + ``[action_tokens×H]`` and the mask is ``[1] + [0]*(H-1)``: the first + action token opens a causal block and the rest attend bidirectionally + within it. + """ + model_dtype = self.action_in_proj.weight.dtype + noisy_actions = noisy_actions.to(dtype=model_dtype) + + time_cond = self.embed_timestep(timestep) + action_emb = self.action_in_proj(noisy_actions) # (B, H, W) + + bsize, action_len = action_emb.shape[:2] + pad_masks = torch.ones(bsize, action_len, dtype=torch.bool, device=action_emb.device) + att_masks = torch.tensor( + [1] + [0] * (self.action_horizon - 1), + dtype=action_emb.dtype, + device=action_emb.device, + )[None, :].expand(bsize, -1) + return action_emb, pad_masks, att_masks, time_cond + + # ── Denoising step ─────────────────────────────────────────────── + def denoise_step( + self, + prefix_pad_masks: torch.Tensor, + past_key_values, + x_t: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + """Apply one flow-matching denoising step: predict ``v_t`` from ``x_t``. + + Signature differs from π0's by exactly one argument: no ``state``. + """ + suffix_embs, suffix_pad_masks, suffix_att_masks, time_cond = self.embed_suffix(x_t, timestep) + + batch_size = prefix_pad_masks.shape[0] + suffix_len = suffix_pad_masks.shape[1] + prefix_len = prefix_pad_masks.shape[1] + + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + + # Position IDs continue from where the prefix's last valid token left off. + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) + + outputs_embeds, _ = self.paligemma_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=False, + adarms_cond=time_cond, + ) + + # Every suffix token is an action token in π0.5 (no state token to drop). + suffix_out = outputs_embeds[1][:, -self.action_horizon :] + suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype) + return self.action_out_proj(suffix_out) + + # ── Full action generation ─────────────────────────────────────── + @torch.no_grad() + def sample_actions( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + noise: torch.Tensor | None = None, + num_steps: int | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> torch.Tensor: + """Generate an action chunk via iterative flow-matching denoising. + + Convention: ``t=1`` is noise, ``t=0`` is the target — opposite of the + published π0 paper but matching both OpenPI and LeRobot. + + Takes no ``state``: π0.5's state rides inside ``lang_tokens``. + """ + if num_steps is None: + num_steps = self.num_inference_steps + + bsize = lang_tokens.shape[0] + device = lang_tokens.device + if noise is None: + noise_shape = (self.action_horizon, self.action_dim) + if isinstance(generator, list): + if len(generator) != bsize: + raise ValueError(f"Expected {bsize} generators, got {len(generator)}.") + noise = torch.stack( + [torch.randn(noise_shape, dtype=torch.float32, device=device, generator=item) for item in generator] + ) + else: + noise = torch.randn( + bsize, + *noise_shape, + dtype=torch.float32, + device=device, + generator=generator, + ) + + # 1. Prefix embeddings + mask building. + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + images, image_masks, lang_tokens, lang_masks + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks) + + # 2. Forward prefix through PaliGemma LM, producing a list[(k, v)] cache. + _, past_key_values = self.paligemma_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=True, + ) + + # 3. Euler-integrated denoising from t=1 down to t=0. + dt = -1.0 / num_steps + x_t = noise + for step in range(num_steps): + t = 1.0 + step * dt + time_tensor = torch.full((bsize,), t, dtype=torch.float32, device=device) + v_t = self.denoise_step( + prefix_pad_masks=prefix_pad_masks, + past_key_values=past_key_values, + x_t=x_t, + timestep=time_tensor, + ) + x_t = x_t + dt * v_t + return x_t + + # ── Weight loading ─────────────────────────────────────────────── + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + *, + strict: bool = True, + ): + """Load and audit a LeRobot π0.5 safetensors checkpoint. + + Same remap rules as π0 (strip the ``model.`` prefix, flatten→nested + PaliGemma submodules, tied ``lm_head`` → ``embed_tokens``, version-robust + SigLIP nesting), plus two π0.5-specific ones: + + * ``action_time_mlp_{in,out}`` → ``time_mlp_{in,out}``: some + checkpoints were exported under the π0 parameter names. + * ``state_proj.*`` is reported, not silently dropped. A π0.5 + checkpoint should not contain it; its presence usually means a π0 + checkpoint was pointed at the π0.5 model class, which would + otherwise run happily with a randomly-initialized action expert. + + The action-expert norms are AdaRMS here, so they expose ``dense.weight`` + / ``dense.bias`` and no plain ``weight``. A checkpoint that carries a + plain expert-norm ``weight`` is a π0-shaped checkpoint; that too is + rejected rather than skipped. ``strict=False`` exists only for focused + remapping unit tests that intentionally provide a partial state dict; + the serving path always uses the strict default. + """ + params_dict = dict(self.named_parameters()) + buffers_dict = dict(self.named_buffers()) + + _PALIGEMMA_SUBMODULES = ("vision_tower", "multi_modal_projector", "language_model") + _EXPERT_PREFIX = "paligemma_with_expert.gemma_expert.model." + + def _remap(name: str) -> str: + # Strip the leading "model." that LeRobot's PI05Policy wrapper adds. + if name.startswith("model."): + name = name[len("model.") :] + + # π0-style timestep MLP names → π0.5 names. + if name.startswith("action_time_mlp_in."): + name = "time_mlp_in." + name[len("action_time_mlp_in.") :] + elif name.startswith("action_time_mlp_out."): + name = "time_mlp_out." + name[len("action_time_mlp_out.") :] + + # Nested PaliGemma layout. + for sub in _PALIGEMMA_SUBMODULES: + flat = f"paligemma_with_expert.paligemma.{sub}." + nested = f"paligemma_with_expert.paligemma.model.{sub}." + if name.startswith(flat) and not name.startswith(nested): + return nested + name[len(flat) :] + + if name == "paligemma_with_expert.paligemma.lm_head.weight": + return "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + + return name + + def _fix_vision_tower(name: str) -> str: + """Reconcile SigLIP nesting across transformers versions (≤5.3 wraps + the encoder in ``vision_tower.vision_model.*``; ≥5.4 flattens it).""" + vt = "paligemma_with_expert.paligemma.model.vision_tower." + if not name.startswith(vt): + return name + rest = name[len(vt) :] + if name in params_dict or name in buffers_dict: + return name + if rest.startswith("vision_model."): + candidate = vt + rest[len("vision_model.") :] + else: + candidate = vt + "vision_model." + rest + return candidate if (candidate in params_dict or candidate in buffers_dict) else name + + loaded = 0 + skipped: list[str] = [] + pi0_shaped: list[str] = [] + filled_params: set = set() + + for name, loaded_weight in weights: + mapped = _fix_vision_tower(_remap(name)) + + # Diagnose π0-shaped keys instead of dropping them quietly. + is_expert_norm_weight = mapped.startswith(_EXPERT_PREFIX) and ( + mapped.endswith("input_layernorm.weight") + or mapped.endswith("post_attention_layernorm.weight") + or mapped == _EXPERT_PREFIX + "norm.weight" + ) + if mapped.startswith("state_proj.") or is_expert_norm_weight: + pi0_shaped.append(mapped) + continue + + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded += 1 + filled_params.add(mapped) + elif mapped in buffers_dict: + buffers_dict[mapped].copy_(loaded_weight) + loaded += 1 + filled_params.add(mapped) + else: + skipped.append(mapped) + + # LeRobot stores PaliGemma's tied text embedding as lm_head.weight; keep + # lm_head filled so the tied-weight state stays consistent. + embed_key = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + lm_head_key = "paligemma_with_expert.paligemma.lm_head.weight" + if embed_key in filled_params and lm_head_key in params_dict and lm_head_key not in filled_params: + params_dict[lm_head_key].data.copy_(params_dict[embed_key].data) + filled_params.add(lm_head_key) + + # Reverse audit: any model param that got no checkpoint tensor at all + # would be running with random init. + missing_params: list[str] = [] + for pname in params_dict: + if pname in filled_params: + continue + if "rotary_emb" in pname or pname.endswith(".inv_freq"): + continue + missing_params.append(pname) + + parts: list[str] = [] + if pi0_shaped: + parts.append( + f"{len(pi0_shaped)} checkpoint key(s) are π0-shaped, not π0.5-shaped (first 5: {pi0_shaped[:5]})" + ) + if skipped: + parts.append(f"{len(skipped)} checkpoint key(s) had no model target (first 5: {skipped[:5]})") + if missing_params: + parts.append(f"{len(missing_params)} model param(s) received no weight (first 5: {missing_params[:5]})") + + if parts and strict: + raise RuntimeError("Incomplete or incompatible π0.5 checkpoint: " + "; ".join(parts)) + if parts: + logger.debug("π0.5 partial test load: %d tensors loaded — %s.", loaded, "; ".join(parts)) + else: + logger.info("π0.5 load_weights: %d tensors loaded, 0 skipped, 0 missing.", loaded) + return filled_params + + +EntryClass = Pi05ForActionPrediction diff --git a/vllm_omni/diffusion/models/pi05/pipeline_pi05.py b/vllm_omni/diffusion/models/pi05/pipeline_pi05.py new file mode 100644 index 000000000..a4d32ed81 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/pipeline_pi05.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 VLA pipeline for vllm-omni. + +Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the +DreamZero contract and the π0 pipeline: the pipeline owns ALL preprocessing. It +reads the raw robot observation from +``req.sampling_params.extra_args["robot_obs"]`` (delivered by the OpenPI +realtime serving layer), builds model inputs, runs flow-matching denoising, and +returns ``DiffusionOutput(output={"actions": ndarray})``. + +π0.5 is stateless across calls (no KV reuse, first-order Markov), so +``session_id`` / ``reset`` from the OpenPI protocol are accepted but ignored. + +The post-processing order is load-bearing and matches LeRobot:: + + unnormalize → absolute actions → to_cpu + +``AbsoluteActionsProcessorStep`` must run *after* unnormalization, because a +relative-action checkpoint's ``norm_stats`` are computed in relative space. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import fields as dataclass_fields + +import numpy as np +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.pi05.config import SUPPORTED_DTYPE_NAMES, Pi05Config +from vllm_omni.diffusion.models.pi05.modeling_pi05 import Pi05ForActionPrediction +from vllm_omni.diffusion.models.pi05.processor_pi05 import Pi05Processor +from vllm_omni.diffusion.models.pi05_pipeline_config import PI05_PIPELINE as PI05_PIPELINE +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = init_logger(__name__) + +# π0.5 pins the PaliGemma tokenizer (LeRobot hardcodes it too). +DEFAULT_PI05_TOKENIZER = "google/paligemma-3b-pt-224" + +# float16 is absent deliberately, not by oversight: nothing here has been +# validated against it. The float32/bfloat16 trade-off is in the deploy config +# and recipes/lerobot/Pi05.md. +SUPPORTED_DTYPES = (torch.float32, torch.bfloat16) + +# The two lists guard different entry points — what the checkpoint declares +# versus the dtype actually cast to — so they must not drift apart. +assert {str(dtype).split(".")[-1] for dtype in SUPPORTED_DTYPES} == set(SUPPORTED_DTYPE_NAMES) + + +def _checkpoint_declared_keys(model_dir: str) -> set[str]: + """The keys the checkpoint's config.json actually carries.""" + path = os.path.join(model_dir, "config.json") + if not os.path.exists(path): + return set() + with open(path, encoding="utf-8") as f: + return set(json.load(f)) + + +def _comparable(value): + """Compare yaml and checkpoint values without tripping on list/tuple.""" + return list(value) if isinstance(value, (list, tuple)) else value + + +def _pi05_post_process(x): + """Module-level identity post-process (picklable across the orchestrator's + multiprocess boundary — a local closure is not).""" + return x + + +def get_pi05_post_process_func(od_config: OmniDiffusionConfig): + """π0.5 returns actions directly; post-processing is identity.""" + del od_config + return _pi05_post_process + + +_LEROBOT_FLOAT32_IN_BFLOAT16 = ( + "vision_tower", + "multi_modal_projector", + "input_layernorm", + "post_attention_layernorm", + "model.norm", +) + + +def _to_bfloat16_for_inference(model: Pi05ForActionPrediction) -> None: + """Apply LeRobot's mixed-precision inference layout explicitly.""" + inner_model = model.paligemma_with_expert + for name, param in inner_model.named_parameters(): + target_dtype = ( + torch.float32 if any(selector in name for selector in _LEROBOT_FLOAT32_IN_BFLOAT16) else torch.bfloat16 + ) + param.data = param.data.to(dtype=target_dtype) + for buffer in inner_model.buffers(): + if buffer.is_floating_point(): + buffer.data = buffer.data.to(dtype=torch.bfloat16) + + for name, param in model.named_parameters(): + if not name.startswith("paligemma_with_expert."): + param.data = param.data.to(dtype=torch.float32) + for name, buffer in model.named_buffers(): + if not name.startswith("paligemma_with_expert.") and buffer.is_floating_point(): + buffer.data = buffer.data.to(dtype=torch.float32) + + +def _set_inference_dtype(model: Pi05ForActionPrediction, dtype: torch.dtype) -> None: + if dtype == torch.float32: + model.to(dtype=torch.float32) + elif dtype == torch.bfloat16: + _to_bfloat16_for_inference(model) + else: + raise ValueError(f"Unsupported π0.5 inference dtype: {dtype!r}.") + + +class Pi05Pipeline(nn.Module): + """π0.5 VLA pipeline: raw robot obs → continuous action chunk. + + Registered as ``"Pi05Pipeline"`` in the diffusion registry. + """ + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.prefix = prefix + self.model_dir = self._resolve_model_dir(od_config.model) + self.config = self._build_config(od_config) + + custom_args = od_config.custom_pipeline_args or {} + self.tokenizer_source = str(custom_args.get("tokenizer", self._resolve_tokenizer_source())) + + self._torch_dtype = self._resolve_dtype(od_config) + self._device = self._resolve_device(od_config) + + self.tokenizer = self._load_tokenizer() + self.model = self._initialize_model() + + self.processor = Pi05Processor(self.config, self.tokenizer, self._device) + + # ------------------------------------------------------------------ + # Construction helpers + # ------------------------------------------------------------------ + @staticmethod + def _resolve_model_dir(model: str | None) -> str | None: + """Return a local directory for ``model``; download an HF repo id if needed.""" + if not model: + return None + if os.path.isdir(model): + return model + # Via repo_utils' shared HfApi rather than huggingface_hub directly, so + # the download carries vLLM-Omni's user agent like every other repo access. + from vllm_omni.transformers_utils.repo_utils import hf_api + + return hf_api().snapshot_download( + repo_id=model, + allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"], + ) + + def _build_config(self, od_config: OmniDiffusionConfig) -> Pi05Config: + """Read the config from the checkpoint, then let the deploy yaml override it.""" + checkpoint_config = Pi05Config.from_pretrained(self.model_dir) if self.model_dir else None + if checkpoint_config is None: + return Pi05Config.from_model_config(od_config.model_config) + if not od_config.model_config: + return checkpoint_config + + resolved = { + item.name: getattr(checkpoint_config, item.name) for item in dataclass_fields(Pi05Config) if item.init + } + # Only a key the checkpoint actually declares can disagree with the yaml. + # Comparing against the dataclass default instead would report every + # serving-only field, such as policy_server_config, on every start. + declared_keys = _checkpoint_declared_keys(self.model_dir) + for key, value in od_config.model_config.items(): + if key in resolved and key in declared_keys and _comparable(value) != _comparable(resolved[key]): + logger.warning( + "Pi05Pipeline: the deploy config sets %s=%r, overriding %r from the checkpoint.", + key, + value, + resolved[key], + ) + resolved.update(od_config.model_config) + return Pi05Config.from_model_config(resolved) + + def _resolve_tokenizer_source(self) -> str: + """Prefer the checkpoint dir if it ships tokenizer files; else PaliGemma.""" + if self.model_dir and os.path.isdir(self.model_dir): + if os.path.exists(os.path.join(self.model_dir, "tokenizer_config.json")): + return self.model_dir + return DEFAULT_PI05_TOKENIZER + + @staticmethod + def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype: + """Resolve the dtype the weights are actually cast to. + + This is the load-bearing check, not ``Pi05Config.dtype``: the cast in + :meth:`_initialize_model` reads the *top-level* ``OmniDiffusionConfig`` + field, so a guard on the model config alone would let an unsupported + dtype through. See :data:`SUPPORTED_DTYPES` for why float16 is excluded. + """ + dt = od_config.dtype + resolved = dt if isinstance(dt, torch.dtype) else getattr(torch, str(dt).split(".")[-1], None) + if resolved not in SUPPORTED_DTYPES: + raise ValueError( + f"Unsupported π0.5 dtype: {dt!r}. Supported: " + f"{', '.join(sorted(str(d).split('.')[-1] for d in SUPPORTED_DTYPES))}." + ) + return resolved + + @staticmethod + def _resolve_device(od_config: OmniDiffusionConfig) -> torch.device: + from vllm_omni.diffusion.distributed.utils import get_local_device + + try: + return get_local_device() + except Exception: # noqa: BLE001 + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def _load_tokenizer(self): + from transformers import AutoTokenizer + + # padding_side="right" is part of the π0.5 spec: the prefix is a fixed + # 200-token block whose live tokens must start at index 0. + return AutoTokenizer.from_pretrained(self.tokenizer_source, padding_side="right") + + def has_real_checkpoint(self) -> bool: + return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors")) + + def _initialize_model(self) -> Pi05ForActionPrediction: + if not self.has_real_checkpoint(): + expected = os.path.join(self.model_dir or "", "model.safetensors") + raise FileNotFoundError(f"π0.5 serving requires checkpoint weights at {expected}.") + model = Pi05ForActionPrediction(self.config) + _set_inference_dtype(model, self._torch_dtype) + model.to(device=self._device) + self._load_checkpoint(model) + model.eval() + return model + + def _load_checkpoint(self, model: Pi05ForActionPrediction) -> None: + import safetensors.torch + + path = os.path.join(self.model_dir, "model.safetensors") + logger.info("Pi05Pipeline: loading π0.5 weights from %s", path) + try: + state = safetensors.torch.load_file(path) + model.load_weights(state.items()) + except Exception as exc: + raise RuntimeError(f"Failed to load complete π0.5 checkpoint {path}: {exc}") from exc + + # ------------------------------------------------------------------ + # Framework weight-loading hook + # ------------------------------------------------------------------ + def load_weights(self, weights=()): # noqa: D401 + """No-op for the diffusion loader: π0.5 self-loads its checkpoint.""" + for _ in weights: + pass + return None + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: + extra_args = getattr(req.sampling_params, "extra_args", None) or {} + robot_obs = extra_args.get("robot_obs") + + if robot_obs is None: + # Dummy warmup path (no obs): return zeros so engine warmup/capture + # doesn't crash. Mirrors DreamZero's dummy-run handling. + first_prompt = req.prompts[0] if req.prompts else "" + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + num_steps = getattr(req.sampling_params, "num_inference_steps", None) + if prompt == "dummy run" or num_steps == 1: + logger.info("Pi05Pipeline: dummy warmup request without robot_obs — returning zeros.") + return DiffusionOutput( + output={ + "actions": np.zeros( + (self.config.chunk_size, self.config.action_dim), + dtype=np.float32, + ) + }, + ) + return DiffusionOutput( + error="Pi05Pipeline.forward requires sampling_params.extra_args['robot_obs'].", + ) + + images, image_masks, lang_tokens, lang_masks = self.processor.build_model_inputs(robot_obs) + + num_steps = getattr(req.sampling_params, "num_inference_steps", None) + if num_steps is not None and ( + isinstance(num_steps, bool) or not isinstance(num_steps, (int, np.integer)) or int(num_steps) < 1 + ): + raise ValueError(f"num_inference_steps must be a positive integer, got {num_steps!r}.") + if num_steps is not None: + num_steps = int(num_steps) + + actions = self.model.sample_actions( + images=images, + image_masks=image_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + num_steps=num_steps, + generator=req.sampling_params.generator, + ) + + return DiffusionOutput(output={"actions": self.processor.build_model_outputs(actions, robot_obs)}) diff --git a/vllm_omni/diffusion/models/pi05/processor_pi05.py b/vllm_omni/diffusion/models/pi05/processor_pi05.py new file mode 100644 index 000000000..11ebad2db --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/processor_pi05.py @@ -0,0 +1,568 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +r"""Preprocessing for the π0.5 VLA model. + +Converts a raw robot observation (multi-camera images + language instruction + +proprioceptive state) into the tensors ``Pi05ForActionPrediction.sample_actions`` +consumes. + +**The defining π0.5 difference:** the state is not projected by a ``state_proj`` +layer. It is normalized to ``[-1, 1]``, discretized into ``state_num_bins`` bins, +and serialized into the language prompt:: + + "Task: , State: ... ;\nAction: " + +so ``sample_actions`` receives no state tensor at all. + +**Normalization must precede discretization.** The discretizer bins over +``[-1, 1]`` and assumes the state is already in that range. Reversed, every bin +index is wrong and nothing raises. + +Reference: + - OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad) + - OpenPI: openpi/src/openpi/models/pi0_config.py (PaliGemmaTokenizer.tokenize) + - LeRobot: lerobot/src/lerobot/policies/pi05/processor_pi05.py +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.pi05.config import resolve_excluded_action_indices +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + DEFAULT_IMAGE_RESOLUTION, + DEFAULT_MAX_TOKEN_LEN, + DEFAULT_STATE_NUM_BINS, +) + +logger = init_logger(__name__) + +# LeRobot NormalizerProcessorStep.eps. +NORM_EPS = 1e-8 + + +# ────────────────────────────────────────────────────────────────────── +# Image preprocessing (identical to π0 — SigLIP is unchanged in π0.5) +# ────────────────────────────────────────────────────────────────────── +def resize_with_pad( + images: torch.Tensor, + target_height: int, + target_width: int, + mode: str = "bilinear", +) -> torch.Tensor: + """Resize ``(B, C, H, W)`` images to the target shape, preserving aspect + ratio with -1 padding on the short side. + + Matches openpi ``image_tools.resize_with_pad_torch`` — the clamp to + [-1, 1] is what lets the padded region blend with SigLIP-normalized + pixels without adding signal at the boundary. + """ + if images.ndim != 4: + raise ValueError(f"Expected 4-D (B,C,H,W), got {images.ndim}-D") + _, _, cur_h, cur_w = images.shape + ratio = max(cur_w / target_width, cur_h / target_height) + rh, rw = int(cur_h / ratio), int(cur_w / ratio) + align_corners = False if mode == "bilinear" else None + resized = F.interpolate(images, size=(rh, rw), mode=mode, align_corners=align_corners) + resized = resized.clamp(-1.0, 1.0) + ph, rem_h = divmod(target_height - rh, 2) + pw, rem_w = divmod(target_width - rw, 2) + return F.pad(resized, (pw, pw + rem_w, ph, ph + rem_h), value=-1.0) + + +def pil_image_to_tensor(image: Image.Image) -> torch.Tensor: + """PIL → ``(1, C, H, W)`` float32 in ``[-1, 1]`` (SigLIP normalization).""" + if image.mode != "RGB": + image = image.convert("RGB") + arr = np.array(image, dtype=np.float32) / 255.0 * 2.0 - 1.0 + return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + + +class Pi05ImageProcessor: + """Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor.""" + + def __init__(self, image_size: int = DEFAULT_IMAGE_RESOLUTION[0]): + self.image_size = image_size + + def preprocess_single(self, image: Any) -> torch.Tensor: + """Convert one LeRobot-domain image to a normalized model tensor. + + PIL and uint8 inputs have domain ``[0, 255]``. Floating ndarray/tensor + inputs must be finite and in ``[0, 1]``. Values are converted + deterministically to ``[-1, 1]``; the pixel content is never used to + guess its input domain. + """ + t = self._to_tensor(image) + if t.shape[2] != self.image_size or t.shape[3] != self.image_size: + t = resize_with_pad(t, self.image_size, self.image_size) + return t + + def _to_tensor(self, image: Any) -> torch.Tensor: + if isinstance(image, Image.Image): + return pil_image_to_tensor(image) + if isinstance(image, np.ndarray): + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"Expected an HWC ndarray with 3 channels, got shape {image.shape}.") + t = torch.from_numpy(np.ascontiguousarray(image)).permute(2, 0, 1).unsqueeze(0) + elif isinstance(image, torch.Tensor): + t = image.unsqueeze(0) if image.ndim == 3 else image + if t.ndim != 4 or t.shape[0] != 1 or t.shape[1] != 3: + raise ValueError(f"Expected a CHW tensor (optionally batched once), got shape {tuple(t.shape)}.") + else: + raise TypeError(f"Unsupported image type for π0.5 preprocessing: {type(image)}") + + if t.dtype == torch.uint8: + return t.to(dtype=torch.float32) / 255.0 * 2.0 - 1.0 + if not t.is_floating_point(): + raise TypeError(f"π0.5 images must use uint8 or floating dtype, got {t.dtype}.") + if not torch.isfinite(t).all(): + raise ValueError("Floating π0.5 image contains NaN or Inf.") + low, high = float(t.amin()), float(t.amax()) + if low < 0.0 or high > 1.0: + raise ValueError(f"Floating π0.5 images must be in [0, 1]; observed min={low}, max={high}.") + return t.to(dtype=torch.float32) * 2.0 - 1.0 + + def make_empty_image(self) -> torch.Tensor: + """Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot.""" + return torch.full((1, 3, self.image_size, self.image_size), -1.0) + + +# ────────────────────────────────────────────────────────────────────── +# Normalization (LeRobot NormalizerProcessorStep) +# ────────────────────────────────────────────────────────────────────── +@dataclass(frozen=True) +class NormStats: + """The affine map a checkpoint's statistics define. + + ``mean_std`` carries mean and std; the two range modes carry the lower and + upper bound (``min``/``max``, or the ``q01``/``q99`` quantiles π0.5 ships). + """ + + mode: str + lower: torch.Tensor + upper: torch.Tensor + + +def build_norm_stats(norm_stats: dict | None, key: str) -> NormStats | None: + """Read ``norm_stats[key]`` into tensors, or ``None`` when it carries none. + + The mode comes from the entry. A LeRobot state_dict ships mean, std, min, + max, q01 and q99 at once, so the statistic names present cannot select it. + Missing modes default to quantile, which is what LeRobot defaults π0.5's + STATE and ACTION to (π0 defaults to mean_std). + """ + entry = (norm_stats or {}).get(key) + if not entry: + return None + + mode = str(entry.get("mode", "quantile")).lower() + if mode == "mean_std": + names = ("mean", "std") + elif mode == "min_max": + names = ("min", "max") + elif mode == "quantile": + mode, names = "min_max", ("q01", "q99") + else: + raise ValueError(f"Unsupported normalization mode for {key!r}: {mode!r}.") + + bounds = [entry.get(name) for name in names] + if any(bound is None for bound in bounds): + raise ValueError(f"Normalization mode {mode!r} for {key!r} requires {names[0]!r} and {names[1]!r}.") + lower, upper = (torch.as_tensor(bound, dtype=torch.float32) for bound in bounds) + return NormStats(mode=mode, lower=lower, upper=upper) + + +def apply_norm(x: torch.Tensor, stats: NormStats | None, *, inverse: bool = False) -> torch.Tensor: + """Map raw units to the model's normalized space, or back with ``inverse``. + + Actions are padded to ``max_action_dim`` while the statistics cover only the + real width, so the tail passes through untouched. The eps rules are + LeRobot's: ``mean_std`` always divides by ``std + eps``, and the range modes + substitute ``eps`` only for an exactly zero range. + """ + if stats is None: + return x + lower = stats.lower.to(device=x.device, dtype=x.dtype) + upper = stats.upper.to(device=x.device, dtype=x.dtype) + valid = lower.shape[0] + head = x[..., :valid] + + if stats.mode == "mean_std": + head = head * upper + lower if inverse else (head - lower) / (upper + NORM_EPS) + else: + span = upper - lower + span = torch.where(span == 0, NORM_EPS, span) + head = (head + 1.0) * 0.5 * span + lower if inverse else 2.0 * (head - lower) / span - 1.0 + + if valid == x.shape[-1]: + return head + return torch.cat([head, x[..., valid:]], dim=-1) + + +# ────────────────────────────────────────────────────────────────────── +# State: normalize → discretize → prompt (π0.5's defining path) +# ────────────────────────────────────────────────────────────────────── +def _pad_or_truncate(raw_state: Any, width: int) -> np.ndarray: + """Zero-pad / truncate a vector to ``(width,)`` float32.""" + if raw_state is None: + return np.zeros((width,), dtype=np.float32) + if isinstance(raw_state, torch.Tensor): + raw_state = raw_state.detach().cpu().numpy() + state = np.asarray(raw_state, dtype=np.float32).reshape(-1) + if state.shape[0] < width: + state = np.pad(state, (0, width - state.shape[0])) + elif state.shape[0] > width: + state = state[:width] + return state.astype(np.float32) + + +def as_state_vector(raw_state: Any, state_dim: int) -> np.ndarray: + """Coerce a request's raw state to ``(state_dim,)`` float32, or raise. + + LeRobot's ``Pi05PrepareStateTokenizerProcessorStep`` discretizes whatever + width it is handed and never pads to ``max_state_dim``, so zero-filling a + 7-dim state up to 32 would append 25 state tokens the checkpoint never saw + in training. + """ + if raw_state is None: + raise ValueError("π0.5 requires a state; there is no default to fall back on.") + if isinstance(raw_state, torch.Tensor): + raw_state = raw_state.detach().cpu().numpy() + state = np.asarray(raw_state, dtype=np.float32) + if state.ndim == 2 and state.shape[0] == 1: + state = state[0] + if state.ndim != 1: + raise ValueError(f"π0.5 state must be shaped (D,) or (1, D); got {tuple(state.shape)}.") + if not np.isfinite(state).all(): + raise ValueError("π0.5 state contains NaN or Inf.") + if state.shape[0] != state_dim: + raise ValueError( + f"π0.5 state has {state.shape[0]} dimension(s), but the checkpoint declares " + f"{state_dim}. The state is serialized into the prompt at its real width, so " + "padding or truncating it would change the tokens the model sees." + ) + return state + + +def discretize_state(state: np.ndarray, *, num_bins: int = DEFAULT_STATE_NUM_BINS) -> np.ndarray: + """Discretize a ``[-1, 1]`` state into ``num_bins`` integer bins. + + Byte-for-byte LeRobot's ``processor_pi05.py``:: + + np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 + + A state below ``-1`` lands in bin ``-1``, and that negative bin is part of + the contract: the checkpoint was trained with ``" -1"`` in the state prompt + for those dimensions. Clipping it to ``0`` changes the tokens the model + sees, so this must not clip. + + ``bins`` stays float64, matching LeRobot's default ``linspace`` dtype, so + boundary values fall on the same side. + + **Assumes the state is already normalized** — see the module docstring. + """ + bins = np.linspace(-1.0, 1.0, num_bins + 1)[:-1] + return (np.digitize(np.asarray(state, dtype=np.float32), bins=bins) - 1).astype(np.int64) + + +def build_pi05_prompt( + *, + task: str, + normalized_state: np.ndarray, + state_num_bins: int = DEFAULT_STATE_NUM_BINS, +) -> str: + """Build the π0.5 prompt: instruction + serialized discretized state. + + Matches LeRobot's ``Pi05PrepareStateTokenizerProcessorStep``, including the + task cleanup (``strip``, ``_`` → space, newline → space), the exact template + and the state values it serializes. The template already ends in a newline, + so — unlike π0 — there is no separate newline-appending step. + + ``normalized_state`` comes from ``Pi05ForActionPrediction._normalize_state``. + """ + cleaned_task = (task or "").strip().replace("_", " ").replace("\n", " ") + bins = discretize_state(normalized_state, num_bins=state_num_bins) + state_str = " ".join(str(int(x)) for x in bins.tolist()) + return f"Task: {cleaned_task}, State: {state_str};\nAction: " + + +def tokenize_prompt(tokenizer, text: str, max_token_len: int = DEFAULT_MAX_TOKEN_LEN): + """Return ``(input_ids, attention_mask)`` lists, length exactly ``max_token_len``. + + ``padding="max_length"`` is what makes the prefix a constant shape: the text + segment is always ``max_token_len`` tokens regardless of the instruction, so + only ``attention_mask.sum()`` varies per request. + """ + enc = tokenizer( + text, + padding="max_length", + max_length=max_token_len, + truncation=True, + add_special_tokens=True, + return_tensors=None, + ) + return list(enc["input_ids"]), list(enc["attention_mask"]) + + +# ────────────────────────────────────────────────────────────────────── +# Relative actions +# ────────────────────────────────────────────────────────────────────── +class Pi05RelativeActions: + """The relative/absolute action transform, as a single paired object. + + LeRobot builds one ``RelativeActionsProcessorStep`` and hands *the same + instance* to ``AbsoluteActionsProcessorStep``; the two directions must + agree on ``enabled`` and on which dimensions are excluded, so they are one + object here too. + + **Deviation from LeRobot, on purpose.** LeRobot's step keeps the reference + state on ``self`` between the pre- and post-pass. That is safe for a + single-threaded training loop and unsafe for a server: two in-flight + requests would share one reference state and silently corrupt each other's + actions. Here the state is passed explicitly to :meth:`to_absolute`, so the + object stays immutable after construction and is safe to share across + requests. + + Transform (LeRobot / OpenPI): ``relative = action - state`` on the way in, + ``absolute = relative + state`` on the way out, applied only to the + dimensions *not* named in ``exclude_joints``. Gripper open/close is an + absolute command, which is why it is excluded by default. + """ + + def __init__( + self, + *, + enabled: bool, + exclude_joints: list[str] | None = None, + action_names: list[str] | None = None, + max_action_dim: int = 32, + ): + self.enabled = bool(enabled) + self.exclude_joints = list(exclude_joints or []) + self.action_names = list(action_names) if action_names else None + self.max_action_dim = int(max_action_dim) + + # Boolean mask over action dims: True = this dim is relative to state. + mask = np.ones((self.max_action_dim,), dtype=bool) + if self.enabled: + for idx in resolve_excluded_action_indices(self.exclude_joints, self.action_names): + if 0 <= idx < self.max_action_dim: + mask[idx] = False + # Padding beyond the real action dimensions carries no signal; + # leaving it "relative" would add state noise into dead channels. + if self.action_names: + mask[len(self.action_names) :] = False + else: + mask[:] = False + self.relative_mask = mask + + @property + def num_relative_dims(self) -> int: + return int(self.relative_mask.sum()) + + def _state_row(self, state: Any, device, dtype) -> torch.Tensor: + """Raw state → ``(B, max_action_dim)`` aligned with the action dims. + + Accepts one state for the whole batch (``(D,)``, the serving path, + where the pipeline runs B=1) or one state per sample (``(B, D)``). + LeRobot caches the batched state and shifts each sample by its own row, + so the per-sample form has to be honoured: ``_pad_or_truncate`` + flattens, which would silently reduce ``(B, D)`` to sample 0's state and + apply it to every sample — wrong answers, no error. + """ + if isinstance(state, torch.Tensor): + state = state.detach().cpu().numpy() + arr = np.asarray(0.0 if state is None else state, dtype=np.float32) + if arr.ndim > 2: + raise ValueError(f"Expected state shaped (D,) or (B, D), got {tuple(arr.shape)}") + rows = arr if arr.ndim == 2 else arr[None, :] + padded = np.stack([_pad_or_truncate(row, self.max_action_dim) for row in rows]) + return torch.as_tensor(padded, device=device, dtype=dtype) + + def to_relative(self, actions: torch.Tensor, state: Any) -> torch.Tensor: + """``absolute → relative``. Input side (step 3). + + Not used on the inference path — there are no input actions to convert + at serving time — but it is what the transform *means*, and the parity + test exercises it as the inverse of :meth:`to_absolute`. + """ + if not self.enabled: + return actions + return self._shift(actions, state, sign=-1.0) + + def to_absolute(self, actions: torch.Tensor, state: Any) -> torch.Tensor: + """``relative → absolute``. Output side (step 2 of the post-pipeline). + + ``state`` must be the **raw** state — the same one the model was given + before normalization — because relative actions live in raw action space. + """ + if not self.enabled: + return actions + return self._shift(actions, state, sign=+1.0) + + def _shift(self, actions: torch.Tensor, state: Any, *, sign: float) -> torch.Tensor: + if actions.ndim != 3: + raise ValueError(f"Expected actions shaped (B, horizon, action_dim), got {tuple(actions.shape)}") + action_dim = actions.shape[-1] + if action_dim != self.max_action_dim: + raise ValueError( + f"Action dim {action_dim} does not match max_action_dim={self.max_action_dim}; " + "the relative mask would be misaligned." + ) + state_row = self._state_row(state, actions.device, actions.dtype) # (1, D) + mask = torch.as_tensor(self.relative_mask, device=actions.device) + delta = torch.where(mask, state_row, torch.zeros_like(state_row)) + # Broadcast over the action horizon: every step of the chunk is + # expressed relative to the same current state. + return actions + sign * delta[:, None, :] + + +# ────────────────────────────────────────────────────────────────────── +# Model input assembly +# ────────────────────────────────────────────────────────────────────── +def _extract_images(robot_obs: dict, config) -> dict[str, Any]: + """Pull a ``{feature_key: image}`` map out of a raw robot obs. + + This is functional step 1 (``rename_observations``): keys are translated + through ``config.image_key_map`` so serving wire names map onto the + checkpoint's ``input_features`` identities. + """ + images = robot_obs.get("images") + if not isinstance(images, dict): + images = {k: v for k, v in robot_obs.items() if _is_image_like(v)} + key_map = config.image_key_map or {} + return {key_map.get(k, k): v for k, v in images.items() if _is_image_like(v)} + + +def _is_image_like(value: Any) -> bool: + if isinstance(value, (Image.Image, torch.Tensor)): + return True + if isinstance(value, np.ndarray): + return value.ndim >= 3 + if isinstance(value, (list, tuple)): + try: + return np.asarray(value).ndim >= 3 + except Exception: # noqa: BLE001 + return False + return False + + +def _assemble_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device, state_norm: NormStats | None): + """Convert a raw robot observation into ``sample_actions`` inputs. + + Returns ``(images, image_masks, lang_tokens, lang_masks)`` — note there is + **no state tensor**: π0.5 carries the state inside ``lang_tokens``. + + Camera slots follow ``config.image_feature_keys`` order. A missing key keeps + its semantic slot and receives an empty image with a false mask. The output + always contains exactly ``max_cameras`` slots for the deployed model. + """ + image_size = int(config.image_resolution[0]) + img_proc = Pi05ImageProcessor(image_size=image_size) + max_cameras = config.max_cameras + + feature_keys = config.image_feature_keys or [] + obs_images = _extract_images(robot_obs, config) + if not feature_keys: + # No declared camera order — fall back to whatever the obs provides, + # preserving insertion order. + feature_keys = list(obs_images.keys())[:max_cameras] + + camera_keys = feature_keys[:max_cameras] + if all(obs_images.get(key) is None for key in camera_keys): + raise ValueError("π0.5 observation must provide at least one configured camera image.") + if not config.image_feature_keys and len(obs_images) > max_cameras: + raise ValueError( + f"π0.5 observation provides {len(obs_images)} cameras, which exceeds max_cameras={max_cameras}." + ) + + images: list[torch.Tensor] = [] + image_masks: list[torch.Tensor] = [] + for key in camera_keys: + image = obs_images.get(key) + if image is None: + tensor = img_proc.make_empty_image().to(device=device) + mask = False + else: + tensor = img_proc.preprocess_single(image).to(device=device) + mask = True + images.append(tensor) + image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device)) + + # A checkpoint may declare fewer named slots than the serving graph. + while len(images) < max_cameras: + images.append(img_proc.make_empty_image().to(device=device)) + image_masks.append(torch.tensor([False], dtype=torch.bool, device=device)) + + # Normalize and serialize the state into the tokenized prompt. + raw_state = as_state_vector(robot_obs.get("state"), config.state_dim) + normalized_state = apply_norm(torch.from_numpy(raw_state), state_norm).numpy() + prompt = build_pi05_prompt( + task=robot_obs.get("prompt", "") or "", + normalized_state=normalized_state, + state_num_bins=config.state_num_bins, + ) + ids, attn = tokenize_prompt(tokenizer, prompt, config.tokenizer_max_length) + lang_tokens = torch.tensor([ids], dtype=torch.long, device=device) + lang_masks = torch.tensor([attn], dtype=torch.bool, device=device) + + return images, image_masks, lang_tokens, lang_masks + + +# ────────────────────────────────────────────────────────────────────── +# The interface the pipeline uses +# ────────────────────────────────────────────────────────────────────── +class Pi05Processor: + """Everything between the OpenPI wire and the model. + + Owns the normalization statistics and the relative-action transform, so the + pipeline holds no preprocessing state of its own. + """ + + def __init__(self, config, tokenizer, device: torch.device): + self.config = config + self.tokenizer = tokenizer + self.device = device + self._state_norm = build_norm_stats(config.norm_stats, "state") + self._action_norm = build_norm_stats(config.norm_stats, "action") + self.relative_actions = Pi05RelativeActions( + enabled=config.use_relative_actions, + exclude_joints=config.relative_exclude_joints, + action_names=config.action_feature_names, + max_action_dim=config.max_action_dim, + ) + if self._action_norm is None: + logger.info("π0.5: no action normalization stats; returned actions stay in normalized space.") + if self.relative_actions.enabled: + logger.info( + "π0.5: relative actions enabled — %d of %d action dims are state-relative (excluded joints: %s).", + self.relative_actions.num_relative_dims, + config.max_action_dim, + config.relative_exclude_joints, + ) + + def build_model_inputs(self, robot_obs: dict): + """Raw observation → ``sample_actions`` inputs. + + No state tensor comes back: π0.5 carries the state inside ``lang_tokens``. + """ + return _assemble_model_inputs(robot_obs, self.config, self.tokenizer, self.device, self._state_norm) + + def build_model_outputs(self, actions: torch.Tensor, robot_obs: dict) -> np.ndarray: + """Model output → the action chunk the robot receives. + + ``AbsoluteActionsProcessorStep`` runs after unnormalization, because a + relative-action checkpoint's statistics are computed in relative space. + The raw state is the one the prompt encoded, before normalization. + """ + actions = apply_norm(actions, self._action_norm, inverse=True) + if self.relative_actions.enabled: + actions = self.relative_actions.to_absolute(actions, robot_obs.get("state")) + return actions.squeeze(0)[..., : self.config.action_dim].float().cpu().numpy() diff --git a/vllm_omni/diffusion/models/pi05_pipeline_config.py b/vllm_omni/diffusion/models/pi05_pipeline_config.py new file mode 100644 index 000000000..27035d509 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05_pipeline_config.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lightweight pipeline topology for π0.5.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +# π0.5 is one diffusion stage (robot observation -> action chunk via flow +# matching). It is registered in ``OMNI_PIPELINES`` because π0.5 is online-served +# via ``vllm serve --deploy-config pi05.yaml``. Keep this topology outside the +# ``pi05`` package: importing a package submodule executes ``pi05.__init__``, +# which imports the runtime pipeline and ``diffusion.data``. +PI05_PIPELINE = PipelineConfig( + model_type="pi05", + model_arch="Pi05Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Pi05Pipeline", + ), + ), +) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 26b052a3a..f94952b1d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -206,6 +206,11 @@ _DIFFUSION_MODELS = { "pipeline_pi0", "Pi0Pipeline", ), + "Pi05Pipeline": ( + "pi05", + "pipeline_pi05", + "Pi05Pipeline", + ), "LongCatImageEditPipeline": ( "longcat_image", "pipeline_longcat_image_edit", @@ -382,6 +387,7 @@ _NO_CACHE_ACCELERATION = { # to it; list it here so a stray cache_backend override disables gracefully # instead of erroring. "Pi0Pipeline", + "Pi05Pipeline", "LingBotWorldCausalDMDPipeline", } @@ -586,6 +592,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "MingImagePipeline": "get_ming_image_post_process_func", "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", "Pi0Pipeline": "get_pi0_post_process_func", + "Pi05Pipeline": "get_pi05_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "FluxKontextPipeline": "get_flux_kontext_post_process_func", diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 11fde2bf8..450cfbb79 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -194,6 +194,9 @@ class ServingRealtimeRobotOpenPI: # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI # protocol fields on top. seed = obs.pop("seed", None) + # Nested so engine knobs cannot collide with robot-defined obs keys. + sampling = obs.get("sampling_params") or {} + robot_obs = {key: value for key, value in obs.items() if key != "sampling_params"} sampling_params = OmniDiffusionSamplingParams() for default_params in get_default_sampling_params_list(self.engine_client): if isinstance(default_params, OmniDiffusionSamplingParams): @@ -205,15 +208,16 @@ class ServingRealtimeRobotOpenPI: { "reset": reset, "session_id": session_id, - "robot_obs": obs, + "robot_obs": robot_obs, } ) prompt = obs.get("prompt", "") - sampling_params = OmniDiffusionSamplingParams( - seed=int(seed) if seed is not None else None, - extra_args=extra_args, - ) + if seed is not None: + sampling_params.seed = int(seed) + if "num_inference_steps" in sampling: + sampling_params.num_inference_steps = sampling["num_inference_steps"] + sampling_params.extra_args = extra_args return OmniDiffusionRequest( prompt=prompt, sampling_params=sampling_params, From 5b22f5562f630ba571cc70d99f104ff01ba87f49 Mon Sep 17 00:00:00 2001 From: Tiago de Souza Fernandes <43066400+Tiagosf00@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:26:34 +0800 Subject: [PATCH 32/99] fix(minicpmo45): deadline-align native duplex silence continuation (#7059) Signed-off-by: Tiagosf00 --- tests/engine/duplex/test_deadline_pacing.py | 88 +++++ .../duplex/test_deadline_pacing_scheduler.py | 370 ++++++++++++++++++ tests/engine/duplex/test_session_manager.py | 4 + vllm_omni/engine/duplex/plugin.py | 5 + .../engine/duplex/session/append_task.py | 6 +- .../engine/duplex/session/model_channel.py | 9 + vllm_omni/engine/duplex/session/runner.py | 108 ++++- .../models/minicpmo_4_5/duplex/session.py | 7 + 8 files changed, 580 insertions(+), 17 deletions(-) create mode 100644 tests/engine/duplex/test_deadline_pacing.py create mode 100644 tests/engine/duplex/test_deadline_pacing_scheduler.py diff --git a/tests/engine/duplex/test_deadline_pacing.py b/tests/engine/duplex/test_deadline_pacing.py new file mode 100644 index 000000000..2b1bf09d6 --- /dev/null +++ b/tests/engine/duplex/test_deadline_pacing.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Deterministic unit tests for deadline-aligned silence continuation. + +The engine-resident duplex scheduler in +``vllm_omni/engine/duplex/session/runner.py`` aligns each silence +continuation to ``submission_time_N + chunk_period`` and sleeps only the +remaining budget. These tests cover the pure deadline arithmetic and the +per-session reset semantics of the shared ``DuplexModelSessionState`` +contract. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.engine.duplex.session.runner import ( + compute_silence_continuation_deadline, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( + MiniCPMO45ServingSessionState, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("chunk_period_s", "now", "last_submit", "current_deadline", "delay_s", "next_silence_deadline"), + [ + # First continuation anchors to the last real submission: unit N was + # submitted at t=1.0, audio N produced at 1.4 -> 0.6 s of sleep left. + (1.0, 1.4, 1.0, None, 0.6, 3.0), + # Only the remaining budget is slept (audio produced at 1.7). + (1.0, 1.7, 1.0, None, 0.3, 3.0), + # Overdue deadline: no sleep; the chain still advances by the period. + (1.0, 2.5, 1.0, None, 0.0, 3.0), + # No submission yet (None): the first continuation anchors to now. + (1.0, 0.4, None, None, 1.0, 2.4), + # Deadlines advance from the current deadline, not from now, so pipeline + # processing time does not accumulate as timer drift. + (1.0, 4.2, 3.0, 5.0, 0.8, 6.0), + # Two consecutive units that overrun their deadline by 0.2 s keep the + # 1 s cadence: no sleep, and the chain advances one period each. + (1.0, 2.2, 1.0, None, 0.0, 3.0), + (1.0, 3.2, 2.0, 3.0, 0.0, 4.0), + # A short overshoot (within one period of the current deadline) still + # chases the stale deadline: immediate submit, chain advances. + (1.0, 2.5, 1.0, 2.0, 0.0, 3.0), + # A long stall (more than one period past the current deadline) submits + # one continuation immediately and restarts from that submission. + (1.0, 5.0, 1.0, 2.0, 0.0, 6.0), + # The same recovery applies to the first continuation after a stall. + (1.0, 5.0, 1.0, None, 0.0, 6.0), + # Zero chunk period (guarded upstream): total, no sleep. + (0.0, 10.0, None, None, 0.0, 10.0), + ], +) +def test_compute_silence_continuation_deadline( + chunk_period_s: float, + now: float, + last_submit: float | None, + current_deadline: float | None, + delay_s: float, + next_silence_deadline: float, +) -> None: + delay, next_dl = compute_silence_continuation_deadline( + chunk_period_s=chunk_period_s, + now=now, + last_submit=last_submit, + current_deadline=current_deadline, + ) + assert delay == pytest.approx(delay_s) + assert next_dl == pytest.approx(next_silence_deadline) + + +class TestSilenceDeadlineSessionState: + def test_clear_continuation_resets_the_deadline_chain(self) -> None: + # The native-duplex session state feeds the shared runner, so it must + # start with an unset chain and drop it at turn boundaries. + state = MiniCPMO45ServingSessionState() + assert state.last_native_submit_monotonic is None + assert state.silence_deadline_monotonic is None + state.last_native_submit_monotonic = 5.0 + state.silence_deadline_monotonic = 6.0 + state.clear_continuation() + assert state.last_native_submit_monotonic is None + assert state.silence_deadline_monotonic is None diff --git a/tests/engine/duplex/test_deadline_pacing_scheduler.py b/tests/engine/duplex/test_deadline_pacing_scheduler.py new file mode 100644 index 000000000..451af0bfd --- /dev/null +++ b/tests/engine/duplex/test_deadline_pacing_scheduler.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Focused scheduler tests for deadline-aligned silence continuation. + +These tests exercise the production ``_schedule_silence_continuation`` path +(and the append acceptance callback it installs) rather than copying the +scheduling arithmetic. The monotonic clock is a fake so the deadline math is +deterministic, and ``asyncio.sleep`` is simulated by advancing that clock. +Append completion is controlled by gating the stage port's ``submit``. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any + +import pytest + +import vllm_omni.engine.duplex.session.model_channel as model_channel_module +import vllm_omni.engine.duplex.session.runner as runner_module +from tests.engine.duplex.test_session_runner import ( + Harness, + append_audio, + close_harness, + open_harness, + tts_output, +) +from vllm_omni.engine.duplex.contracts import DuplexStageSubmission + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +CHUNK_PERIOD_S = 1.0 + + +class FakeClock: + """A monotonic clock whose value the test controls directly.""" + + def __init__(self, start: float = 0.0) -> None: + self.value = start + + def __call__(self) -> float: + return self.value + + +async def _active_response_harness() -> Harness: + """Open a harness and drive a TTS segment so a response is active.""" + h = await open_harness() + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + assert h.session.active_response_id is not None + return h + + +def _continuation_kwargs(h: Harness) -> dict[str, object]: + """The scheduler arguments the model channel would pass for this session.""" + return { + "request_id": h.session.active_request_id, + "owner_id": f"response:{h.session.active_response_id}", + "response_id": h.session.active_response_id, + "response_owned": True, + "expected_epoch": h.session.epoch, + "expected_model_turn_id": h.session.turn_id, + } + + +def _install_fake_clock( + monkeypatch: pytest.MonkeyPatch, + *, + clock: FakeClock, + real_sleep: Callable[..., Any], +) -> list[float]: + """Point both modules at the fake clock and simulate ``asyncio.sleep``. + + ``real_sleep`` is the unpatched ``asyncio.sleep`` captured before patching; + the simulated sleep advances the fake clock by the requested delay and then + yields control so the append task can run. The returned ``stall`` list is + added once to the next *positive* sleep so a test can model a wake that + happens late (after the planned deadline); zero-duration scheduling yields + do not consume it. + """ + monkeypatch.setattr(runner_module.time, "monotonic", clock) + monkeypatch.setattr(model_channel_module.time, "monotonic", clock) + stall: list[float] = [0.0] + + async def fake_sleep(delay: float) -> None: + clock.value += delay + if delay > 0: + clock.value += stall[0] + stall[0] = 0.0 + await real_sleep(0) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + return stall + + +@pytest.mark.asyncio +async def test_normal_cadence_keeps_deadlines_aligned(monkeypatch: pytest.MonkeyPatch) -> None: + """Processing time reduces the remaining sleep; consecutive deadlines align. + + Each accepted unit becomes the new anchor (its submission time), and the + following continuation advances from the stored deadline rather than + resetting to now. + """ + clock = FakeClock(start=100.0) + real_sleep = asyncio.sleep + _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + try: + # Real input accepted at t=100.0; the model processes for 0.4 s. + h.runner.model_state.last_native_submit_monotonic = 100.0 + clock.value = 100.4 + + # First continuation: due at 101.0, so only 0.6 s of sleep remains. + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + first_task = h.runner.tasks.append_tail + assert first_task is not None + assert await first_task + # The accepted unit re-anchored the chain to its own submission time + # (101.0) and stored the following deadline (102.0). + assert h.runner.model_state.last_native_submit_monotonic == pytest.approx(101.0) + assert h.runner.model_state.silence_deadline_monotonic == pytest.approx(102.0) + + # Model produces the continuation's audio; the next unit is planned + # from the stored deadline (no drift): the following deadline is 103.0. + clock.value = 101.3 + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + second_task = h.runner.tasks.append_tail + assert second_task is not None + assert await second_task + assert h.runner.model_state.silence_deadline_monotonic == pytest.approx(103.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_in_flight_real_append_skips_outdated_silence(monkeypatch: pytest.MonkeyPatch) -> None: + """A real append accepted after the silence was planned re-anchors the chain. + + The planned silence queues behind the real append; once the real append + completes and re-anchors, ``before_append`` detects the changed anchor + (numeric timestamp) and the outdated silence is skipped (no submission, + no callback overwrite). + """ + clock = FakeClock(start=200.0) + real_sleep = asyncio.sleep + _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + state = h.runner.model_state + submissions_before = len(h.port.submissions) + + original_submit = h.port.submit + entered = asyncio.Event() + release = asyncio.Event() + + async def gated_submit(submission: DuplexStageSubmission) -> Any: + entered.set() + await release.wait() + return await original_submit(submission) + + monkeypatch.setattr(h.port, "submit", gated_submit) + try: + # Existing anchor at t=200.0. + state.last_native_submit_monotonic = 200.0 + state.silence_deadline_monotonic = None + + # The real append starts at t+0.8 and blocks inside the gate. + clock.value = 200.8 + real_append_task = await h.runner._start_append( + h.runner.model.silence_unit_payload(), + final=False, + silence_continuation=False, + ) + await entered.wait() + assert not real_append_task.done() + + # At t+1.0 the silence is planned using the existing anchor; its append + # task waits behind the real append. + clock.value = 201.0 + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + silence_task = h.runner.tasks.append_tail + assert silence_task is not None + assert not silence_task.done() + + # At t+1.1 the real append is released; both appends finish. + clock.value = 201.1 + release.set() + assert await real_append_task + assert await silence_task + + # Exactly one additional submission happened: the real append. The + # outdated silence was skipped, so its callback never overwrote the + # new anchor (the real append's submission time t+0.8), and the + # deadline the real append cleared stays cleared. + assert len(h.port.submissions) == submissions_before + 1 + assert state.last_native_submit_monotonic == pytest.approx(200.8) + assert state.silence_deadline_monotonic is None + + # A fresh continuation uses the new anchor. The clock is left at t+1.1: + # the scheduler sleeps the remaining 0.7 s (deadline t+1.8), which + # verifies re-anchoring controls the wait. Following deadline t+2.8. + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + fresh_task = h.runner.tasks.append_tail + assert fresh_task is not None + assert await fresh_task + assert state.last_native_submit_monotonic == pytest.approx(201.8) + assert state.silence_deadline_monotonic == pytest.approx(202.8) + finally: + release.set() + await close_harness(h) + + +@pytest.mark.asyncio +async def test_long_stall_saves_future_deadline_from_actual_submission( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A submission more than one period past the planned deadline restarts. + + The schedule plans the following deadline at ``planned + period``, but the + actual submission happens much later. The accepted continuation must save + ``submit_time + period`` so the next unit is not immediately due. + """ + clock = FakeClock(start=300.0) + real_sleep = asyncio.sleep + stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + # Plan a continuation whose deadline is 1.0 s in the future. The wake + # is delayed 2 s past the planned submission (anchor + 3.0 during the + # scheduler's sleep), so the append accepts 2 s late. + anchor = clock.value + state.last_native_submit_monotonic = anchor + state.silence_deadline_monotonic = None + stall[0] = 2.0 + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + append_task = h.runner.tasks.append_tail + assert append_task is not None + assert await append_task + # anchor + 3.0 > anchor + 2.0 (planned deadline + period) -> save + # submit_time + period = anchor + 4.0. + assert state.last_native_submit_monotonic == pytest.approx(anchor + 3.0) + assert state.silence_deadline_monotonic == pytest.approx(anchor + 4.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_small_wakeup_delay_preserves_planned_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ordinary wakeup jitter (within one period) keeps the planned cadence. + + A submission 0.2 s past the planned deadline must not reset the chain: the + stored deadline stays at ``planned + period`` so drift does not accumulate. + """ + clock = FakeClock(start=400.0) + real_sleep = asyncio.sleep + stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + anchor = clock.value + state.last_native_submit_monotonic = anchor + state.silence_deadline_monotonic = None + stall[0] = 0.2 + + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + append_task = h.runner.tasks.append_tail + assert append_task is not None + assert await append_task + + # Submitted at anchor + period + 0.2 (within one period): the planned + # following deadline is kept. + assert state.last_native_submit_monotonic == pytest.approx(anchor + CHUNK_PERIOD_S + 0.2) + assert state.silence_deadline_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_exact_one_period_late_preserves_planned_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exactly one period late is not "more than one period late". + + The strict ``>`` comparison means a submission exactly at the planned + following deadline keeps the planned cadence (no extra catch-up unit). + """ + clock = FakeClock(start=500.0) + real_sleep = asyncio.sleep + stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + anchor = clock.value + state.last_native_submit_monotonic = anchor + state.silence_deadline_monotonic = None + stall[0] = 1.0 + + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + append_task = h.runner.tasks.append_tail + assert append_task is not None + assert await append_task + + # Submitted exactly at the planned following deadline: kept. + assert state.last_native_submit_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S) + assert state.silence_deadline_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_failed_append_leaves_timing_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed append never advances the timing chain.""" + clock = FakeClock(start=600.0) + real_sleep = asyncio.sleep + _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + state.last_native_submit_monotonic = 600.0 + state.silence_deadline_monotonic = 601.0 + before_anchor = state.last_native_submit_monotonic + before_deadline = state.silence_deadline_monotonic + h.port.fail_submit = RuntimeError("boom") + + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + ) + assert scheduled is True + append_task = h.runner.tasks.append_tail + assert append_task is not None + # The append fails: its task reports failure, and timing is untouched. + assert await append_task is False + assert state.last_native_submit_monotonic == before_anchor + assert state.silence_deadline_monotonic == before_deadline + finally: + await close_harness(h) diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py index 0323f16d8..4f1962ee5 100644 --- a/tests/engine/duplex/test_session_manager.py +++ b/tests/engine/duplex/test_session_manager.py @@ -120,6 +120,8 @@ class FakeModelSessionState(DuplexModelSessionState): self.continuation_units = 0 self.pending_silence_task = None self.pending_silence_owner_id = None + self.last_native_submit_monotonic = None + self.silence_deadline_monotonic = None def retain_committed_audio(self, payload, *, operation_id, reserved_bytes=0) -> None: self.committed_audio_payload = payload @@ -136,6 +138,8 @@ class FakeModelSessionState(DuplexModelSessionState): def clear_continuation(self) -> None: self.continuation_owner_id = None self.continuation_units = 0 + self.last_native_submit_monotonic = None + self.silence_deadline_monotonic = None class FakeDataPlane(DuplexDataPlane): diff --git a/vllm_omni/engine/duplex/plugin.py b/vllm_omni/engine/duplex/plugin.py index 130d5204d..f9f9be283 100644 --- a/vllm_omni/engine/duplex/plugin.py +++ b/vllm_omni/engine/duplex/plugin.py @@ -123,6 +123,11 @@ class DuplexModelSessionState(ABC): continuation_units: int pending_silence_task: asyncio.Task[bool] | None pending_silence_owner_id: str | None + # Deadline-aligned silence continuation state: the monotonic submission + # time of the most recent native input unit and the next silence + # continuation deadline. A real (non-silence) input resets the chain. + last_native_submit_monotonic: float | None + silence_deadline_monotonic: float | None @abstractmethod def retain_committed_audio( diff --git a/vllm_omni/engine/duplex/session/append_task.py b/vllm_omni/engine/duplex/session/append_task.py index 724dc7f82..2bc060d63 100644 --- a/vllm_omni/engine/duplex/session/append_task.py +++ b/vllm_omni/engine/duplex/session/append_task.py @@ -53,8 +53,11 @@ class AppendAttempt: retained_committed_payload: dict[str, object] | None #: Response reserved before submission, to be failed if submission does not happen. precreated_response_id: str | None + #: Commits timing state once the runtime accepts the append (before any + #: returned output event can clear the continuation chain). + on_append_accepted: Callable[[float], None] | None = None #: Last chance to call the append off, checked once the predecessor is done. - before_append: Callable[[], bool] | None + before_append: Callable[[], bool] | None = None # ------------------------------------------------------------------ # # Compensation # @@ -140,6 +143,7 @@ class AppendAttempt: ), final=self.final, expected_epoch=self.epoch, + on_append_accepted=self.on_append_accepted, ) if append_ok: model_state.context_locked = True diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py index 322db7ede..13f0f7781 100644 --- a/vllm_omni/engine/duplex/session/model_channel.py +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -22,6 +22,7 @@ tracked task, and aborting a stage request in the background. from __future__ import annotations import asyncio +import time import uuid from collections.abc import Awaitable, Callable, Mapping from dataclasses import replace @@ -146,12 +147,16 @@ class ModelChannel: operation_id: str | None = None, final: bool, expected_epoch: int | None = None, + on_append_accepted: Callable[[float], None] | None = None, ) -> tuple[bool, bool]: session = self._ctx.session if not session.capabilities.supports_input_append: return True, False if expected_epoch is not None and session.epoch != expected_epoch: return True, False + # Anchor the submission time before the RPC; the acceptance callback + # commits timing state only if the append actually submitted. + submit_time = time.monotonic() try: result = await self._append_via_data_plane( payload, @@ -169,6 +174,10 @@ class ModelChannel: return True, False if expected_epoch is not None and session.epoch != expected_epoch: return True, False + # Commit timing state before returned output events can clear the + # silence-continuation chain (e.g. a terminal turn-end). + if on_append_accepted is not None: + on_append_accepted(submit_time) request_id, _ = duplex_data_plane_request_info(result) if request_id is not None: self._ctx.plugin.data_plane.begin_request(request_id) diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index 0b8af49ab..ef464a5fa 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -21,6 +21,7 @@ orchestrator loop (the session is never touched from another thread): from __future__ import annotations import asyncio +import time import uuid from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field @@ -126,6 +127,38 @@ class _Internal: _CANCEL_EVENTS = frozenset({"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}) +def compute_silence_continuation_deadline( + *, + chunk_period_s: float, + now: float, + last_submit: float | None, + current_deadline: float | None, +) -> tuple[float, float]: + """Return the silence continuation schedule ``(delay_s, next_silence_deadline)``. + + The first continuation anchors to the latest accepted append's submission + time (``last_submit + chunk_period_s``), falling back to ``now`` when no + submission exists. Later continuations advance from the current deadline + instead of from ``now``, so pipeline processing time does not accumulate + as timer drift. When the schedule is more than one period overdue it is + stale: one continuation submits immediately and the schedule restarts + from ``now`` (``next_silence_deadline = now + chunk_period_s``) instead of + firing a burst of catch-ups. + + ``delay_s`` is the wait before this continuation and + ``next_silence_deadline`` the deadline for the following one. + """ + if current_deadline is None: + base = last_submit if last_submit is not None else now + deadline = base + chunk_period_s + else: + deadline = current_deadline + if now - deadline > chunk_period_s: + deadline = now + delay_s = max(0.0, deadline - now) + return delay_s, deadline + chunk_period_s + + class DuplexSessionRunner: """Owns one ``DuplexEngineSession`` on the orchestrator loop (see module docstring).""" @@ -918,12 +951,22 @@ class DuplexSessionRunner: operation_id: str | None = None, retained_committed_payload: dict[str, object] | None = None, silence_continuation: bool = False, + on_append_accepted: Callable[[float], None] | None = None, before_append: Callable[[], bool] | None = None, ) -> asyncio.Task[bool]: session = self.session model_state = self.model_state if not silence_continuation: self._mark_pending_silence_superseded() + if on_append_accepted is None: + # A real (non-silence) input re-anchors the silence pacing + # chain: its submission becomes the anchor and the stored + # deadline is cleared until the next continuation sets one. + def _reanchor_chain(submit_time: float) -> None: + model_state.last_native_submit_monotonic = submit_time + model_state.silence_deadline_monotonic = None + + on_append_accepted = _reanchor_chain append_epoch = session.epoch append_turn_id = payload_turn_id(payload) if append_turn_id is None: @@ -949,6 +992,7 @@ class DuplexSessionRunner: operation_id=operation_id, retained_committed_payload=retained_committed_payload, precreated_response_id=session.active_response_id if precreate_response else None, + on_append_accepted=on_append_accepted, before_append=before_append, ) @@ -1032,29 +1076,60 @@ class DuplexSessionRunner: append_tail = self.tasks.append_tail if (append_tail is None or append_tail.done()) and self._real_input_waiting(): return False - continuation_delay_s = max(0.0, float(session.capabilities.chunk_period_ms or 1000) / 1000.0) - if continuation_delay_s > 0: - await asyncio.sleep(continuation_delay_s) - if ( - self.tasks.append_tail is not append_tail - or ((append_tail is None or append_tail.done()) and self._real_input_waiting()) - or self.model.silence_continuation_is_stale( + chunk_period_s = max(0.0, float(session.capabilities.chunk_period_ms or 1000) / 1000.0) + # Snapshot the anchor after any wait for pending silence. A real append + # accepted between the snapshot and this continuation's submission + # re-anchors the chain; _still_valid() then skips the stale unit. + anchor = model_state.last_native_submit_monotonic + # Align the next silence unit to submission_time_N + chunk_period and + # sleep only the remaining budget. The deadline is stored by the + # acceptance callback when the append actually submits, so skipped or + # stale continuations never advance the clock. + delay_s, next_silence_deadline = compute_silence_continuation_deadline( + chunk_period_s=chunk_period_s, + now=time.monotonic(), + last_submit=anchor, + current_deadline=model_state.silence_deadline_monotonic, + ) + if delay_s > 0: + await asyncio.sleep(delay_s) + if ( + self.tasks.append_tail is not append_tail + or ((append_tail is None or append_tail.done()) and self._real_input_waiting()) + or self.model.silence_continuation_is_stale( + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + ): + return False + + def _still_valid() -> bool: + return ( + # The anchor changed (a real append was accepted) after this + # continuation was planned; the unit is outdated. + model_state.last_native_submit_monotonic == anchor + and not self._real_input_waiting() + and not self.model.silence_continuation_is_stale( request_id=request_id, response_id=response_id, response_owned=response_owned, expected_epoch=expected_epoch, expected_model_turn_id=expected_model_turn_id, ) - ): - return False + ) - def _still_valid() -> bool: - return not self._real_input_waiting() and not self.model.silence_continuation_is_stale( - request_id=request_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_model_turn_id=expected_model_turn_id, + def _on_append_accepted(submit_time: float) -> None: + # Commit timing state once the runtime accepts the append. If the + # submission is more than one chunk period past the planned + # deadline, the schedule is stale: restart from the actual + # submission. Small delays keep the planned cadence so ordinary + # jitter does not accumulate as drift. + model_state.last_native_submit_monotonic = submit_time + model_state.silence_deadline_monotonic = ( + submit_time + chunk_period_s if submit_time > next_silence_deadline else next_silence_deadline ) model_state.pending_silence_owner_id = owner_id @@ -1062,6 +1137,7 @@ class DuplexSessionRunner: dict(payload) if isinstance(payload, dict) else {}, final=False, silence_continuation=True, + on_append_accepted=_on_append_accepted, before_append=_still_valid, ) return task is not None diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py index 49d0caab6..fe53ec653 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py @@ -29,6 +29,11 @@ class MiniCPMO45ServingSessionState(DuplexModelSessionState): continuation_units: int = 0 pending_silence_task: asyncio.Task[bool] | None = None pending_silence_owner_id: str | None = None + # Deadline-aligned silence continuation state: the monotonic submission + # time of the most recent native input unit and the next silence + # continuation deadline. A real (non-silence) input resets the chain. + last_native_submit_monotonic: float | None = None + silence_deadline_monotonic: float | None = None def retain_committed_audio( self, @@ -55,3 +60,5 @@ class MiniCPMO45ServingSessionState(DuplexModelSessionState): self.continuation_units = 0 self.pending_silence_task = None self.pending_silence_owner_id = None + self.last_native_submit_monotonic = None + self.silence_deadline_monotonic = None From c61b30255bc2099a947ccf7bf0955f73c139d277 Mon Sep 17 00:00:00 2001 From: "howard.shan" Date: Thu, 17 Sep 2026 20:47:10 +0800 Subject: [PATCH 33/99] [Qwen3-Omni] Remove redundant text embedding in DeepStack path (#7573) Signed-off-by: howard-shan Signed-off-by: howard.shan --- benchmarks/qwen3_omni/embedding_harness.py | 226 ++++++++++++++++++ .../qwen3_omni/test_qwen3_omni_embed.py | 214 +++++++++++++++++ .../qwen3_omni/qwen3_omni_moe_thinker.py | 57 ++--- 3 files changed, 470 insertions(+), 27 deletions(-) create mode 100644 benchmarks/qwen3_omni/embedding_harness.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py diff --git a/benchmarks/qwen3_omni/embedding_harness.py b/benchmarks/qwen3_omni/embedding_harness.py new file mode 100644 index 000000000..63d14ef0d --- /dev/null +++ b/benchmarks/qwen3_omni/embedding_harness.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Small fixtures for the real Qwen3 thinker embedding method and its oracle.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from types import SimpleNamespace + +import torch +from torch import nn +from torch.nn import functional as F + +SCENARIOS = ("text", "empty", "audio", "image", "video", "mixed", "interleaved", "vision_no_deepstack") + + +@dataclass(frozen=True) +class ModelShape: + vocab_size: int = 152064 + hidden_size: int = 2048 + deepstack_levels: int = 3 + image_token_id: int = 151655 + video_token_id: int = 151656 + audio_token_id: int = 151675 + + def __post_init__(self): + if min(self.vocab_size, self.hidden_size, self.deepstack_levels) <= 0: + raise ValueError("vocab_size, hidden_size and deepstack_levels must be positive") + ids = (self.image_token_id, self.video_token_id, self.audio_token_id) + if len(set(ids)) != 3 or not all(0 < token < self.vocab_size for token in ids): + raise ValueError("Provide three distinct, positive, in-vocabulary modality IDs; use oov=True to test OOV") + + def token_ids(self, oov=False): + if oov: + return dict(zip(("image", "video", "audio"), range(self.vocab_size, self.vocab_size + 3))) + return {name: getattr(self, f"{name}_token_id") for name in ("image", "video", "audio")} + + +@dataclass +class Feature: + modality: str + positions: torch.Tensor + main: torch.Tensor + deepstack: torch.Tensor | None + combined: torch.Tensor + + +@dataclass +class EmbeddingCase: + name: str + input_ids: torch.Tensor + is_multimodal: torch.Tensor | None + features: tuple[Feature, ...] + fingerprint: str + + def fresh_embeddings(self): + # The thinker replaces list entries with main-scale views. + return None if self.name == "text" else [feature.combined for feature in self.features] + + +def _tensor_bytes(tensor): + return tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + + +def build_case( + name, + num_tokens, + num_mm_tokens, + shape, + *, + device="cuda", + dtype=torch.bfloat16, + mask_device="cpu", + seed=42, + oov=False, +): + """Build known feature placements; no model merge is used to define them.""" + if name not in SCENARIOS: + raise ValueError(f"Unknown scenario: {name}") + if num_tokens <= 0 or not 0 <= num_mm_tokens < num_tokens: + raise ValueError("Require N > 0 and 0 <= M < N") + if (name in {"text", "empty"}) != (num_mm_tokens == 0): + raise ValueError("Only text/empty scenarios accept M=0") + if name in {"mixed", "interleaved"} and (num_mm_tokens < 4 or num_mm_tokens % 2): + raise ValueError("Mixed/interleaved scenarios require an even M >= 4") + token_ids = shape.token_ids(oov) + generator = torch.Generator().manual_seed(seed) + ids = torch.randint(0, min(shape.vocab_size, *token_ids.values()), (num_tokens,), generator=generator) + mask = torch.zeros(num_tokens, dtype=torch.bool) + positions = torch.arange((num_tokens - num_mm_tokens) // 2, (num_tokens - num_mm_tokens) // 2 + num_mm_tokens) + groups = [] + if name == "mixed": + groups = [("audio", positions[: num_mm_tokens // 2]), ("image", positions[num_mm_tokens // 2 :])] + elif name == "interleaved": + groups = [("video", positions[::2]), ("audio", positions[1::2])] + elif num_mm_tokens: + groups = [("image" if name == "vision_no_deepstack" else name, positions)] + features = [] + for modality, indexes in groups: + ids[indexes] = token_ids[modality] + mask[indexes] = True + rows = torch.arange(len(indexes), dtype=torch.float32)[:, None].remainder(31) / 32 + cols = torch.arange(shape.hidden_size, dtype=torch.float32)[None, :].remainder(7) / 16 + main_cpu = rows + cols + {"audio": 2, "image": 4, "video": 6}[modality] + main = main_cpu.to(device=device, dtype=dtype) + deepstack = None + parts = [main] + if modality != "audio" and name != "vision_no_deepstack": + deepstack = torch.stack( + [(main_cpu + 16 + 4 * level).to(device=device, dtype=dtype) for level in range(shape.deepstack_levels)] + ) + parts.extend(deepstack.unbind(0)) + combined = torch.cat(parts, dim=-1) + # set_mm_embedding_modality simply attaches this attribute. Tag only + # after cat/to: views and copies do not preserve arbitrary attributes. + combined.modality = modality + features.append(Feature(modality, indexes, main, deepstack, combined)) + digest = hashlib.sha256(json.dumps({"name": name, "shape": asdict(shape), "seed": seed, "oov": oov}).encode()) + digest.update(_tensor_bytes(ids)) + digest.update(_tensor_bytes(mask)) + for feature in features: + digest.update(feature.modality.encode()) + digest.update(_tensor_bytes(feature.combined)) + return EmbeddingCase( + name, + ids.to(device=device), + None if name == "text" else mask.to(device=mask_device), + tuple(features), + digest.hexdigest(), + ) + + +class LanguageModelAdapter(nn.Module): + def __init__(self, embedding): + super().__init__() + self.embedding = embedding + + def embed_input_ids(self, input_ids): + return self.embedding(input_ids) + + def forward(self, input_ids): + return self.embed_input_ids(input_ids) + + +def make_thinker( + shape, + *, + device="cuda", + dtype=torch.bfloat16, + seed=42, + backend="torch", + deepstack=True, + buffer_capacity=32, + oov=False, +): + """Skip checkpoint construction, preserving the real class and its MRO.""" + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + ) + + if buffer_capacity <= 0: + raise ValueError("buffer_capacity must be positive") + device = torch.device(device) + if backend == "torch": + embedding = nn.Embedding(shape.vocab_size, shape.hidden_size, device=device, dtype=dtype) + elif backend == "vllm": + from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + with torch.device(device): + embedding = VocabParallelEmbedding(shape.vocab_size, shape.hidden_size, params_dtype=dtype, disable_tp=True) + else: + raise ValueError(f"Unknown embedding backend: {backend}") + with torch.no_grad(): + embedding.weight.uniform_(-0.5, 0.5, generator=torch.Generator(device=device).manual_seed(seed)) + thinker = object.__new__(Qwen3OmniMoeThinkerForConditionalGeneration) + nn.Module.__init__(thinker) + thinker.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=shape.hidden_size)) + for modality, token_id in shape.token_ids(oov).items(): + setattr(thinker.config, f"{modality}_token_id", token_id) + thinker.visual = SimpleNamespace( + deepstack_visual_indexes=list(range(shape.deepstack_levels)) if deepstack else None + ) + thinker.language_model = LanguageModelAdapter(embedding) + thinker._has_oov_mm_tokens = oov + thinker.use_deepstack = deepstack + thinker.deepstack_num_level = shape.deepstack_levels if deepstack else 0 + thinker.deepstack_input_embeds = [ + torch.zeros(buffer_capacity, shape.hidden_size, device=device, dtype=dtype) + for _ in range(thinker.deepstack_num_level) + ] + thinker.deepstack_input_embeds_num_tokens = 0 + return thinker.eval() + + +@torch.inference_mode() +def expected_outputs(model, case): + """Independent oracle: direct lookup and explicit positions, not merge helpers.""" + ids = case.input_ids + if model._has_oov_mm_tokens and case.is_multimodal is not None: + ids = ids.masked_fill(case.is_multimodal.to(ids.device), 0) + main = F.embedding(ids, model.language_model.embedding.weight) + has_deepstack = any(feature.deepstack is not None for feature in case.features) + deepstack = main.new_zeros((model.deepstack_num_level, *main.shape)) if has_deepstack else None + for feature in case.features: + indexes = feature.positions.to(main.device) + main[indexes] = feature.main + if feature.deepstack is not None: + deepstack[:, indexes, :] = feature.deepstack + return main, deepstack + + +@torch.inference_mode() +def assert_case_output(model, case, result): + expected, deepstack = expected_outputs(model, case) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + if deepstack is None: + assert model.deepstack_input_embeds_num_tokens == 0, "Unexpected stale DeepStack buffer" + else: + assert model.deepstack_input_embeds_num_tokens == case.input_ids.numel() + actual = torch.stack([buffer[: case.input_ids.numel()] for buffer in model.deepstack_input_embeds]) + torch.testing.assert_close(actual, deepstack, rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py new file mode 100644 index 000000000..d041eace0 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Run Qwen3-Omni embedding regressions on a Linux/CUDA environment. + +The oracle assigns features to token positions independently of the model's +merge helpers. LM call counts additionally catch issue #7451: ordinary +multimodal inputs must not repeat the text embedding lookup. Small synthetic +weights suffice; no checkpoint is downloaded or loaded. +""" + +import sys + +import pytest +import torch +from pytest_mock import MockerFixture + +from benchmarks.qwen3_omni.embedding_harness import ( + ModelShape, + assert_case_output, + build_case, + expected_outputs, + make_thinker, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +_SHAPE = ModelShape( + vocab_size=64, + hidden_size=8, + deepstack_levels=3, + image_token_id=60, + video_token_id=61, + audio_token_id=62, +) + +_SCENARIOS = ( + "text", + "empty", + "audio", + "image", + "video", + "mixed", + "interleaved", + "vision_no_deepstack", +) + + +@pytest.fixture(scope="module", autouse=True) +def require_cuda_environment(): + if sys.platform != "linux" or torch.version.cuda is None or not torch.cuda.is_available(): + pytest.fail( + "This validation suite requires a Linux/CUDA environment; " + "CPU execution and an all-skipped result are not valid substitutes.", + pytrace=False, + ) + if not torch.cuda.is_bf16_supported(including_emulation=False): + pytest.fail("This validation suite requires native BF16 support on the current CUDA device.", pytrace=False) + + +@torch.inference_mode() +def _embed(model, case): + return model.embed_input_ids( + case.input_ids, + multimodal_embeddings=case.fresh_embeddings(), + is_multimodal=case.is_multimodal, + ) + + +@pytest.mark.parametrize( + ("scenario", "dtype"), + [pytest.param(scenario, torch.bfloat16, id=f"{scenario}-bf16") for scenario in _SCENARIOS] + + [pytest.param(scenario, torch.float32, id=f"{scenario}-fp32") for scenario in ("image", "mixed", "interleaved")], +) +def test_embedding_positions_and_single_lm_lookup(mocker: MockerFixture, scenario: str, dtype: torch.dtype): + model = make_thinker(_SHAPE, device="cuda", dtype=dtype, deepstack=scenario != "vision_no_deepstack") + case = build_case( + scenario, + 16, + 0 if scenario in {"text", "empty"} else 8, + _SHAPE, + device="cuda", + dtype=dtype, + mask_device="cpu", + ) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + # Unlike upstream-only mocks, make_thinker returns an instance of Omni's + # actual thinker class, so its zero-argument super() follows the real MRO. + lm_embed.assert_called_once() + assert_case_output(model, case, result) + + +@pytest.mark.parametrize("scenario", ["image", "mixed", "interleaved"]) +def test_repeated_calls_keep_original_deepstack_features(mocker: MockerFixture, scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + original_embeddings = case.fresh_embeddings() + assert original_embeddings + original_shapes = [tensor.shape for tensor in original_embeddings] + original_modalities = [getattr(tensor, "modality", None) for tensor in original_embeddings] + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + previous_list = None + + for _ in range(3): + embeddings = case.fresh_embeddings() + assert embeddings is not original_embeddings + assert embeddings is not previous_list + assert [tensor.shape for tensor in embeddings] == original_shapes + assert [getattr(tensor, "modality", None) for tensor in embeddings] == original_modalities + lm_embed.reset_mock() + + with torch.inference_mode(): + result = model.embed_input_ids( + case.input_ids, + multimodal_embeddings=embeddings, + is_multimodal=case.is_multimodal, + ) + + lm_embed.assert_called_once() + assert_case_output(model, case, result) + # torch.split replaces entries in the supplied list with main-scale + # views. The next call must receive the original multiscale tensors. + assert [tensor.shape for tensor in original_embeddings] == original_shapes + model._clear_deepstack_input_embeds(case.input_ids.numel()) + previous_list = embeddings + + +@pytest.mark.parametrize("scenario", ["audio", "image", "mixed", "interleaved"]) +def test_oov_multimodal_ids_are_masked_before_lm_lookup(mocker: MockerFixture, scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, oov=True) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu", oov=True) + original_ids = case.input_ids.clone() + mask = case.is_multimodal.to(device=case.input_ids.device) + assert torch.all(case.input_ids[mask] >= _SHAPE.vocab_size) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + lm_embed.assert_called_once() + received_ids = lm_embed.call_args.args[0] + torch.testing.assert_close(received_ids, original_ids.masked_fill(mask, 0), rtol=0, atol=0) + torch.testing.assert_close(case.input_ids, original_ids, rtol=0, atol=0) + assert_case_output(model, case, result) + + +@pytest.mark.parametrize("next_scenario", ["text", "empty"]) +def test_padded_deepstack_read_and_clear_do_not_leak_into_decode(mocker: MockerFixture, next_scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, buffer_capacity=32) + image_case = build_case("image", 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + image_result = _embed(model, image_case) + assert_case_output(model, image_case, image_result) + _, expected_deepstack = expected_outputs(model, image_case) + assert expected_deepstack is not None + + num_valid = image_case.input_ids.numel() + num_padded = 24 + # Padding may contain old buffer contents. _get must zero the tail while + # retaining all valid visual features, as the real forward requires. + for buffer in model.deepstack_input_embeds: + buffer[num_valid:num_padded].fill_(77) + padded = model._get_deepstack_input_embeds(num_padded) + assert padded is not None + for level in range(_SHAPE.deepstack_levels): + actual = padded[f"deepstack_input_embeds_{level}"] + torch.testing.assert_close(actual[:num_valid], expected_deepstack[level]) + assert torch.count_nonzero(actual[num_valid:]).item() == 0 + + model._clear_deepstack_input_embeds(num_padded) + assert model.deepstack_input_embeds_num_tokens == 0 + for buffer in model.deepstack_input_embeds: + assert torch.count_nonzero(buffer[:num_padded]).item() == 0 + + decode_case = build_case(next_scenario, 1, 0, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + decode_result = _embed(model, decode_case) + lm_embed.assert_called_once() + assert_case_output(model, decode_case, decode_result) + decode_deepstack = model._get_deepstack_input_embeds(2) + assert decode_deepstack is not None + for level in range(_SHAPE.deepstack_levels): + assert torch.count_nonzero(decode_deepstack[f"deepstack_input_embeds_{level}"]).item() == 0 + + +@torch.inference_mode() +def test_deepstack_set_resizes_for_larger_scheduled_windows(mocker: MockerFixture): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, buffer_capacity=4) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + for num_tokens in (16, 24): + case = build_case("mixed", num_tokens, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + lm_embed.reset_mock() + result = _embed(model, case) + + lm_embed.assert_called_once() + assert_case_output(model, case, result) + assert model.deepstack_input_embeds_num_tokens == num_tokens + assert all(buffer.size(0) >= num_tokens for buffer in model.deepstack_input_embeds) + model._clear_deepstack_input_embeds(num_tokens) + + +@pytest.mark.parametrize("mask_device", ["cpu", "cuda"]) +@pytest.mark.parametrize("scenario", ["image", "interleaved"]) +def test_cuda_embeddings_accept_cpu_or_cuda_mask(mocker: MockerFixture, scenario: str, mask_device: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device=mask_device) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + lm_embed.assert_called_once() + assert result.device.type == "cuda" + assert_case_output(model, case, result) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index c78679c55..5d5edf0df 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -1273,14 +1273,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - inputs_embeds = self._embed_text_input_ids( - input_ids, - self.language_model.embed_input_ids, - is_multimodal=is_multimodal, - ) - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds + return self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) # Detect interleaved audio-in-video early, since it affects # both the deepstack path and the final embedding merge. @@ -1295,12 +1293,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( is_interleaved = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio) - deepstack_input_embeds = None # split the feat dim to obtain multi-scale visual feature has_vision_embeddings = [ embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings ] - if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings): + has_deepstack_embeddings = self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings) + if has_deepstack_embeddings: multiscale_len = len(self.visual.deepstack_visual_indexes) multimodal_embeddings_multiscale = [] @@ -1343,6 +1341,28 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( if not is_interleaved: mm_position_idx += num_tokens + if is_interleaved: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + inputs_embeds = merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_mm_device, + ) + else: + # multimodal_embeddings now contains the main-scale features. + inputs_embeds = super().embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + if has_deepstack_embeddings: deepstack_input_embeds = inputs_embeds.new_zeros( inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1) ) @@ -1358,24 +1378,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) self._set_deepstack_input_embeds(deepstack_input_embeds) - if is_interleaved: - return merge_interleaved_embeddings( - inputs_embeds, - multimodal_embeddings, - is_video, - is_audio, - is_mm_device, - ) - - # Default: standard merge (no interleaving), same as parent class. - # multimodal_embeddings may have been updated above (deepstack - # main-scale). Use super() to stay consistent with the parent - # implementation and avoid issues seen in Qwen2.5-Omni (#34506). - return super().embed_input_ids( - input_ids, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) + return inputs_embeds def forward( self, From eb1abbadf832c625a9560cb6e0142d2cec4998ce Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 17 Sep 2026 20:51:26 +0800 Subject: [PATCH 34/99] [Feature][TTS] MOSS-TTS: place the reference-audio encoder on the code2wav stage's GPU (#7724) Signed-off-by: Canlin Guo --- .../openai_api/test_tts_adapter.py | 109 ++++++++++++++++++ .../openai/tts_adapters/moss_tts.py | 96 ++++++++++++--- 2 files changed, 192 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 500feab35..a2507a056 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -255,6 +255,115 @@ def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker): assert request.input == "Target." +def _moss_adapter_with_stage_configs(stage_configs, mocker): + engine_client = SimpleNamespace( + model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"), + stage_configs=stage_configs, + ) + return MossTTSAdapter(SpeechServingContext(server=mocker.Mock(), engine_client=engine_client)) + + +def _moss_cuda_available(mocker, device_count=8): + mocker.patch("torch.cuda.is_available", return_value=True) + mocker.patch("torch.accelerator.device_count", return_value=device_count) + + +@pytest.mark.parametrize( + "stage_configs,expected", + [ + # Follows the code2wav stage's first device. + ( + [ + SimpleNamespace(model_stage="moss_tts_local", runtime_config=SimpleNamespace(devices="0")), + SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3")), + ], + "cuda:3", + ), + # No codec-named stage: anchors on the last stage of the pipeline. + ( + [ + SimpleNamespace(model_stage="talker", runtime_config=SimpleNamespace(devices="1")), + SimpleNamespace(model_stage="decoder", runtime_config=SimpleNamespace(devices="2")), + ], + "cuda:2", + ), + # Multi-device codec stage pins replica 0 on the first device. + ( + [ + SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="5,6")), + ], + "cuda:5", + ), + # No explicit pinning anywhere: stage worker defaults to cuda:0. + ( + [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices=None))], + "cuda:0", + ), + ], +) +def test_moss_ref_encoder_device_follows_codec_stage(stage_configs, expected, mocker): + import torch + + _moss_cuda_available(mocker) + adapter = _moss_adapter_with_stage_configs(stage_configs, mocker) + assert adapter._resolve_ref_encoder_device() == torch.device(expected) + + +def test_moss_ref_encoder_device_accepts_dict_runtime(mocker): + import torch + + _moss_cuda_available(mocker) + adapter = _moss_adapter_with_stage_configs( + [SimpleNamespace(model_stage="moss_tts_local_codec", runtime={"devices": "2"})], + mocker, + ) + assert adapter._resolve_ref_encoder_device() == torch.device("cuda:2") + + +def test_moss_ref_encoder_device_cpu_when_cuda_unavailable(mocker): + import torch + + mocker.patch("torch.cuda.is_available", return_value=False) + adapter = _moss_adapter_with_stage_configs( + [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3"))], + mocker, + ) + assert adapter._resolve_ref_encoder_device() == torch.device("cpu") + + +def test_moss_ref_encoder_device_cpu_when_index_out_of_range(mocker): + import torch + + _moss_cuda_available(mocker, device_count=2) + adapter = _moss_adapter_with_stage_configs( + [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="5"))], + mocker, + ) + assert adapter._resolve_ref_encoder_device() == torch.device("cpu") + + +def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker): + import torch + + _moss_cuda_available(mocker) + adapter = _moss_adapter_with_stage_configs( + [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3"))], + mocker, + ) + audio_tokenizer = mocker.Mock() + audio_tokenizer.to.return_value = audio_tokenizer + processor = SimpleNamespace(audio_tokenizer=audio_tokenizer) + from_pretrained = mocker.patch("transformers.AutoProcessor.from_pretrained", return_value=processor) + + assert adapter._get_moss_processor() is processor + audio_tokenizer.to.assert_called_once_with(torch.device("cuda", 3)) + audio_tokenizer.eval.assert_called_once_with() + # Lazy cache: no re-load / re-placement on the second call. + assert adapter._get_moss_processor() is processor + from_pretrained.assert_called_once() + audio_tokenizer.to.assert_called_once() + + def test_qwen3_tts_metadata(): assert Qwen3TTSAdapter.backend == "ar" assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 555fec42d..5f5d8ecb9 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast from transformers import AutoModel, AutoTokenizer from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.inputs import tokens_input +from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( @@ -27,6 +28,8 @@ from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realt if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +logger = init_logger(__name__) + class _MossTTSAdapterBase(ARTTSAdapter): accumulate_nonstreaming: bool = False @@ -143,13 +146,81 @@ class _MossTTSAdapterBase(ARTTSAdapter): return "voice_generator" return "tts" + def _codec_stage_devices(self) -> str | None: + """Read the runtime ``devices`` entry of this pipeline's codec stage. + + The codec (code2wav) stage is the final stage of every MOSS pipeline + (``moss_tts_local_codec`` / ``moss_tts_codec``); prefer a stage whose + ``model_stage`` mentions "codec", else anchor on the last stage. The + runtime mapping accepts the resolved ``runtime_config`` object, the + omegaconf ``runtime`` mapping, and the legacy ``yaml_runtime`` dict. + """ + stages = list(getattr(self.engine_client, "stage_configs", None) or ()) + if not stages: + return None + codec_stages = [s for s in stages if "codec" in str(getattr(s, "model_stage", "") or "")] + anchor = codec_stages[-1] if codec_stages else stages[-1] + for attr in ("runtime_config", "runtime", "yaml_runtime"): + runtime_cfg = getattr(anchor, attr, None) + if runtime_cfg is None: + continue + devices = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + if devices is not None and str(devices).strip(): + return str(devices) + devices = getattr(anchor, "devices", None) + return str(devices) if devices is not None and str(devices).strip() else None + + def _resolve_ref_encoder_device(self): + """Pick the device of the API-process reference-audio encoder. + + The encoder follows the code2wav stage onto its first GPU: the stage + worker remaps its ``devices`` entry through ``CUDA_VISIBLE_DEVICES`` + and sees it as logical index 0, while replica launch restores the + parent environment, so the same index names the same physical GPU in + this process. When the stage pins no devices the stage worker and the + encoder both default to the first visible GPU. Falls back to CPU when + CUDA is unavailable or the stage device is not parseable. + """ + import torch # local to avoid pulling torch at module import time + + if not torch.cuda.is_available(): + return torch.device("cpu") + devices_str = self._codec_stage_devices() + if devices_str is None: + # No explicit pinning: the stage worker lands on the first visible + # GPU (vLLM picks cuda:0 of the un-remapped environment). + return torch.device("cuda", 0) + first = devices_str.split(",")[0].strip() + if not first: + return torch.device("cuda", 0) + if first == "cpu": + return torch.device("cpu") + if not first.isdigit(): + logger.warning("MOSS ref encoder: codec stage devices %r not parseable; keeping CPU", devices_str) + return torch.device("cpu") + index = int(first) + device_count = torch.accelerator.device_count() + if index >= device_count: + logger.warning( + "MOSS ref encoder: codec stage device cuda:%d out of range (%d visible); keeping CPU", + index, + device_count, + ) + return torch.device("cpu") + return torch.device("cuda", index) + def _get_moss_processor(self): """Lazily load the upstream MOSS-TTS processor once per server. Cached on ``self._moss_processor_cache``. The processor owns its own - audio_tokenizer (~1.6 B params); we keep it on CPU so it doesn't - compete with the talker (~8 GiB) and codec (~7 GiB) for our 96 GiB - GPU — per-request ref-audio encoding is fast enough on CPU. + audio_tokenizer (~1.6 B params) used for per-request reference-audio + encoding; it is placed on the code2wav stage's GPU (see + ``_resolve_ref_encoder_device``) so cold encodes run on device instead + of paying a CPU forward per request. Note this memory is charged to + the API process, outside each stage's ``gpu_memory_utilization`` + budget — keep some headroom on the codec GPU. """ cached = getattr(self, "_moss_processor_cache", None) if cached is not None: @@ -159,7 +230,9 @@ class _MossTTSAdapterBase(ARTTSAdapter): model_id = self.engine_client.model_config.model proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) if hasattr(proc, "audio_tokenizer"): - proc.audio_tokenizer = proc.audio_tokenizer.to("cpu").eval() + device = self._resolve_ref_encoder_device() + proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval() + logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device) self._moss_processor_cache = proc return proc @@ -188,16 +261,13 @@ class _MossTTSAdapterBase(ARTTSAdapter): ), ) ) - codec = ( - AutoModel.from_pretrained( - codec_path, - trust_remote_code=True, - ) - # Reference encoding runs in the API process. Keep the codec on CPU - # so it does not reserve accelerator memory outside the model workers. - .to("cpu") - .eval() + codec = AutoModel.from_pretrained( + codec_path, + trust_remote_code=True, ) + # Reference encoding runs in the API process; follow the codec + # (code2wav) stage's GPU like the processor path above. + codec = codec.to(self._resolve_ref_encoder_device()).eval() cached = (tokenizer, processor, codec) self._moss_realtime_components = cached return cached From 882bf8bc95b1a9a08af667a7edd006ebd21e4b57 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Thu, 17 Sep 2026 23:32:31 +0800 Subject: [PATCH 35/99] [Frontend] Add a shared realtime web UI for MiniCPM-o and Qwen3-Omni (#7222) (#7585) Signed-off-by: amy-why-3459 --- docs/serving/realtime_duplex_api.md | 5 +- examples/online_serving/minicpmo/README.md | 9 + .../minicpmo/realtime_web/server.py | 190 +----- examples/online_serving/qwen3_omni/README.md | 71 +- .../qwen3_omni/openai_realtime_client.py | 3 +- .../qwen3_omni/realtime_web/__init__.py | 4 + .../qwen3_omni/realtime_web/__main__.py | 7 + .../online_serving/realtime_web/README.md | 221 +++++++ .../online_serving/realtime_web/__init__.py | 4 + .../online_serving/realtime_web/__main__.py | 7 + .../realtime_web/app/index.html | 23 +- .../realtime_web/app/static/app.js | 436 ++++++++---- .../realtime_web/app/static/pcm_worklet.js | 0 .../app/static/playback_worklet.js | 7 + .../app/static/profiles/common.js | 56 ++ .../app/static/profiles/minicpm_native.js | 55 ++ .../app/static/profiles/qwen3_turn.js | 100 +++ .../realtime_web/app/static/styles.css | 3 + .../online_serving/realtime_web/server.py | 229 +++++++ tests/config/test_environment_variables.py | 4 +- tests/engine/duplex/test_engine_session.py | 4 +- tests/engine/duplex/test_session_manager.py | 6 +- tests/engine/duplex/test_vad_backend.py | 31 +- tests/engine/test_duplex_orchestrator.py | 45 ++ .../openai_api/test_api_server_guards.py | 16 +- .../openai_api/test_duplex_api_server.py | 28 + .../test_minicpmo_realtime_web_server.py | 48 +- .../test_minicpmo_realtime_web_static.py | 71 +- tests/examples/test_realtime_web_profiles.cjs | 412 ++++++++++++ tests/examples/test_realtime_web_profiles.py | 24 + .../models/qwen3_omni/test_duplex_plugin.py | 625 ++++++++++++++++++ .../config/environment_variable_inventory.py | 1 + vllm_omni/config/stage_config.py | 2 + vllm_omni/deploy/qwen3_omni_duplex.yaml | 8 + vllm_omni/engine/duplex/config.py | 65 +- vllm_omni/engine/duplex/plugin.py | 25 +- vllm_omni/engine/duplex/realtime_events.py | 117 ++-- .../engine/duplex/session/append_task.py | 9 +- vllm_omni/engine/duplex/session/control.py | 54 ++ .../engine/duplex/session/engine_session.py | 162 +++-- vllm_omni/engine/duplex/session/helpers.py | 33 +- vllm_omni/engine/duplex/session/manager.py | 11 +- .../engine/duplex/session/model_channel.py | 63 +- vllm_omni/engine/duplex/session/runner.py | 99 ++- vllm_omni/engine/duplex/vad.py | 10 +- vllm_omni/engine/duplex_orchestrator.py | 53 +- vllm_omni/entrypoints/duplex/warmup.py | 5 +- vllm_omni/entrypoints/openai/api_server.py | 52 +- .../models/qwen3_omni/duplex/__init__.py | 2 + .../models/qwen3_omni/duplex/input.py | 103 +++ .../models/qwen3_omni/duplex/plugin.py | 390 +++++++++++ .../models/qwen3_omni/duplex/session.py | 59 ++ .../models/qwen3_omni/pipeline.py | 4 +- 53 files changed, 3438 insertions(+), 633 deletions(-) create mode 100644 examples/online_serving/qwen3_omni/realtime_web/__init__.py create mode 100644 examples/online_serving/qwen3_omni/realtime_web/__main__.py create mode 100644 examples/online_serving/realtime_web/README.md create mode 100644 examples/online_serving/realtime_web/__init__.py create mode 100644 examples/online_serving/realtime_web/__main__.py rename examples/online_serving/{minicpmo => }/realtime_web/app/index.html (76%) rename examples/online_serving/{minicpmo => }/realtime_web/app/static/app.js (59%) rename examples/online_serving/{minicpmo => }/realtime_web/app/static/pcm_worklet.js (100%) rename examples/online_serving/{minicpmo => }/realtime_web/app/static/playback_worklet.js (95%) create mode 100644 examples/online_serving/realtime_web/app/static/profiles/common.js create mode 100644 examples/online_serving/realtime_web/app/static/profiles/minicpm_native.js create mode 100644 examples/online_serving/realtime_web/app/static/profiles/qwen3_turn.js rename examples/online_serving/{minicpmo => }/realtime_web/app/static/styles.css (98%) create mode 100644 examples/online_serving/realtime_web/server.py create mode 100644 tests/examples/test_realtime_web_profiles.cjs create mode 100644 tests/examples/test_realtime_web_profiles.py create mode 100644 tests/model_executor/models/qwen3_omni/test_duplex_plugin.py create mode 100644 vllm_omni/deploy/qwen3_omni_duplex.yaml create mode 100644 vllm_omni/model_executor/models/qwen3_omni/duplex/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/duplex/input.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/duplex/plugin.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/duplex/session.py diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index c18b416c7..a20750090 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -768,7 +768,10 @@ events with `server_event_seq > 40` are replayed in order) ``` with camera frames (omni video; rides the append that closes a 1 s model -unit; each entry is a **bare** base64 JPEG/PNG — no `data:` URL prefix) +unit; each entry is a **bare** base64 JPEG/PNG — no `data:` URL prefix). +Only a model whose Stage 0 interleaves a frame track at unit boundaries takes +this field; a turn model has no such boundary and rejects it, expecting +`conversation.item.create` with `input_image` content instead. ```json { diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 67b268a77..ffd50167f 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -316,3 +316,12 @@ them unless `--allow-invalid-clock` is explicit. [`examples/offline_inference/minicpmo/`](../../offline_inference/minicpmo/) - Recipe: [`recipes/OpenBMB/MiniCPM-o-4_5.md`](../../../recipes/OpenBMB/MiniCPM-o-4_5.md) + +### Shared realtime UI implementation + +The `python -m examples.online_serving.minicpmo.realtime_web` command remains +available with the same MiniCPM defaults and required `--ref-audio`. Its assets +now live in [the shared realtime UI](../realtime_web/README.md), with a dedicated +`minicpm-native` profile preserving native duplex input, playback ACKs and camera +frames. Qwen3 uses a separate profile in the same shell, supporting manual STT +turns or Server VAD turns with speech interruption. diff --git a/examples/online_serving/minicpmo/realtime_web/server.py b/examples/online_serving/minicpmo/realtime_web/server.py index ee78fff74..1ef12d186 100644 --- a/examples/online_serving/minicpmo/realtime_web/server.py +++ b/examples/online_serving/minicpmo/realtime_web/server.py @@ -1,185 +1,23 @@ -"""Static host and same-origin Realtime proxy for the full-duplex demo.""" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from __future__ import annotations +"""Backward-compatible MiniCPM entry point for the shared realtime UI.""" -import argparse -import asyncio -import base64 -import contextlib -import hashlib -import json -import logging -from pathlib import Path -from urllib.parse import urlencode +from examples.online_serving.realtime_web.server import ( + APP_DIR, + STATIC_DIR, + _join_ws_url, + build_app, +) +from examples.online_serving.realtime_web.server import ( + main as _main, +) -import uvicorn -import websockets -from fastapi import FastAPI, WebSocket, WebSocketDisconnect -from fastapi.responses import HTMLResponse, Response -from fastapi.staticfiles import StaticFiles - -logger = logging.getLogger(__name__) -APP_DIR = Path(__file__).parent / "app" -STATIC_DIR = APP_DIR / "static" - - -def _join_ws_url(base: str, path: str, query: str) -> str: - return base.rstrip("/") + path + (("?" + query) if query else "") - - -async def _pump_client_to_backend(client: WebSocket, backend) -> None: - try: - while True: - message = await client.receive() - if message["type"] == "websocket.disconnect": - await backend.close() - return - if message.get("text") is not None: - await backend.send(message["text"]) - elif message.get("bytes") is not None: - await backend.send(message["bytes"]) - except (WebSocketDisconnect, websockets.ConnectionClosed): - with contextlib.suppress(Exception): - await backend.close() - - -async def _pump_backend_to_client(client: WebSocket, backend) -> None: - try: - async for message in backend: - if isinstance(message, bytes): - await client.send_bytes(message) - else: - await client.send_text(message) - except (WebSocketDisconnect, websockets.ConnectionClosed): - return - except RuntimeError as exc: - if "websocket.send" in str(exc) and "websocket.close" in str(exc): - return - raise - - -def _expected_proxy_close(exc: BaseException) -> bool: - if isinstance(exc, (WebSocketDisconnect, websockets.ConnectionClosed, asyncio.CancelledError)): - return True - return isinstance(exc, RuntimeError) and "websocket.send" in str(exc) and "websocket.close" in str(exc) - - -def build_app( - *, - ws_backend: str = "ws://127.0.0.1:8099", - model: str = "openbmb/MiniCPM-o-4_5", - public_realtime_url: str | None = None, - ref_audio: str | None = None, -) -> FastAPI: - app = FastAPI(title="Experimental Full-Duplex Web Demo") - index_path = APP_DIR / "index.html" - app_version_hash = hashlib.sha256() - for asset_path in ( - STATIC_DIR / "app.js", - STATIC_DIR / "pcm_worklet.js", - STATIC_DIR / "playback_worklet.js", - ): - app_version_hash.update(asset_path.read_bytes()) - app_version = app_version_hash.hexdigest()[:12] - - ref_audio_uri: str | None = None - if ref_audio: - ref_path = Path(ref_audio) - if not ref_path.is_file(): - raise SystemExit(f"--ref-audio not found: {ref_audio}") - encoded = base64.b64encode(ref_path.read_bytes()).decode("ascii") - ref_audio_uri = f"data:audio/wav;base64,{encoded}" - - @app.get("/", response_class=HTMLResponse) - def index() -> HTMLResponse: - config = json.dumps( - { - "model": model, - "realtimePath": public_realtime_url or "v1/realtime", - "refAudio": ref_audio_uri, - "appVersion": app_version, - }, - ensure_ascii=True, - ) - html = ( - index_path.read_text(encoding="utf-8") - .replace("__FULL_DUPLEX_CONFIG__", config) - .replace("__FULL_DUPLEX_APP_VERSION__", app_version) - ) - return HTMLResponse(html, headers={"Cache-Control": "no-store"}) - - @app.get("/healthz") - def healthz() -> Response: - return Response(content="ok", media_type="text/plain") - - @app.websocket("/v1/realtime") - async def realtime_proxy(websocket: WebSocket) -> None: - await websocket.accept() - query = urlencode(websocket.query_params.multi_items()) - backend_url = _join_ws_url(ws_backend, "/v1/realtime", query) - logger.info("Proxying Realtime WebSocket to %s", backend_url) - try: - async with websockets.connect( - backend_url, - max_size=64 * 1024 * 1024, - ) as backend: - tasks = { - asyncio.create_task(_pump_client_to_backend(websocket, backend)), - asyncio.create_task(_pump_backend_to_client(websocket, backend)), - } - done, pending = await asyncio.wait( - tasks, - return_when=asyncio.FIRST_COMPLETED, - ) - for task in pending: - task.cancel() - await asyncio.gather(*pending, return_exceptions=True) - for result in await asyncio.gather(*done, return_exceptions=True): - if isinstance(result, BaseException) and not _expected_proxy_close(result): - raise result - except (WebSocketDisconnect, websockets.ConnectionClosed): - return - except Exception: - logger.exception("Realtime WebSocket proxy failed") - with contextlib.suppress(Exception): - await websocket.close(code=1011) - - app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") - return app +__all__ = ["APP_DIR", "STATIC_DIR", "_join_ws_url", "build_app", "main"] def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument("--port", type=int, default=7862) - parser.add_argument("--ws-backend", default="ws://127.0.0.1:8099") - parser.add_argument( - "--public-realtime-url", - help="Browser-visible ws:// or wss:// Realtime URL; defaults to the same-origin proxy.", - ) - parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument( - "--ref-audio", - required=True, - help=( - "Reference voice wav for TTS voice cloning, e.g. the " - "official MiniCPM-o-Demo assets/ref_audio/ref_minicpm_signature.wav" - ), - ) - args = parser.parse_args() - - logging.basicConfig(level=logging.INFO) - uvicorn.run( - build_app( - ws_backend=args.ws_backend, - model=args.model, - public_realtime_url=args.public_realtime_url, - ref_audio=args.ref_audio, - ), - host=args.host, - port=args.port, - log_level="info", - ) + _main(default_profile="minicpm-native") if __name__ == "__main__": diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 8d6014db2..5eb9a2042 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -320,24 +320,17 @@ python openai_realtime_client.py \ Server VAD uses the pinned Silero v6.2 ONNX artifact from `istupakov/silero-vad-onnx` and never downloads model files while processing a session. Make the artifact available in the Hugging Face cache before startup, or configure a local artifact in the deployment YAML: ```yaml +base_config: /path/to/vllm-omni/vllm_omni/deploy/qwen3_omni_duplex.yaml duplex_session: server_vad_model_path: /models/silero_vad.onnx ``` -Start from the bundled Qwen3-Omni deployment YAML, add the fields above at the top level, and serve that complete configuration: - -```bash -cp vllm_omni/deploy/qwen3_omni_moe.yaml /path/to/qwen3_omni_server_vad.yaml -# Edit /path/to/qwen3_omni_server_vad.yaml and add duplex_session. -vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ - --omni \ - --port 8091 \ - --deploy-config /path/to/qwen3_omni_server_vad.yaml -``` - -Keep Qwen's default `session_mode: turn`; `duplex_session` enables the Realtime handler without changing scheduler -semantics. Bare `/v1/realtime` selects that handler. The client uses `?duplex=0` only for the existing non-Server-VAD -wire flow; `?duplex=1` remains a supported compatibility alias. +For the engine-owned plugin, use an overlay based on +`vllm_omni/deploy/qwen3_omni_duplex.yaml`, with `duplex_session.server_vad_model_path` +pointing to your local artifact. See the [complete duplex setup](../realtime_web/README.md#qwen3-with-server-vad-automatic-turns). +The default `qwen3_omni_moe.yaml` remains a turn deployment. Duplex mode serves +`/v1/realtime?duplex=1` and `/v1/chat/completions`; the legacy STT handler requires +a turn deployment. The Python client supports the following command-line arguments: @@ -584,3 +577,53 @@ The gradio script supports the following arguments: - `--ip`: Host/IP for Gradio server (default: 127.0.0.1) - `--port`: Port for Gradio server (default: 7861) - `--share`: Share the Gradio demo publicly (creates a public link) + +## Browser voice call (shared realtime UI) + +Run these commands from the repository root with the vLLM-Omni environment activated. + +### Without VAD: manually submit each turn + +Start the backend: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 +``` + +In another terminal, check backend readiness and start the UI: + +```bash +curl --fail http://127.0.0.1:8091/health +python -m examples.online_serving.qwen3_omni.realtime_web \ + --backend ws://127.0.0.1:8091 --stt --port 7863 +``` + +Open `http://localhost:7863`, start a session, speak, and press **Send turn**. +`--stt` is the default. Each turn uses a separate connection without shared model +history; the conversation shown in the browser is a local log. + +### With VAD: automatically submit after silence + +Start the backend with a Server VAD deployment overlay and a compatible Silero +ONNX artifact, following the +[complete VAD setup](../realtime_web/README.md#qwen3-with-server-vad-automatic-turns). +After its health check succeeds, start the UI: + +```bash +python -m examples.online_serving.qwen3_omni.realtime_web \ + --backend ws://127.0.0.1:8091 --vad --port 7863 +``` + +Speak and pause for 500 ms to submit automatically. `--vad` configures the client; +it does not load a VAD model or change the backend deployment by itself. + +VAD mode keeps uploading audio during replies and supports speech interruption. +The engine-owned Qwen duplex plugin maintains conversation history and tracks +played replies through playback acknowledgements. Each committed utterance starts +a new generation request; the model does not use native streaming-input KV decoding. +STT mode pauses microphone upload while the reply is generated and played. +In VAD mode, click **Camera** to upload sampled frames with your spoken question; +see [camera setup and limits](../realtime_web/README.md#qwen-vad-camera-input). +STT remains audio-only. To switch back to manual turns, use the default turn +deployment, restart the UI with `--stt`, refresh the page, and reconnect. Stop the existing UI before reusing port 7863. +See the [shared UI guide](../realtime_web/README.md) for HTTPS access and testing. diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 0a98d69a0..2a6bca0fe 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -130,7 +130,7 @@ async def run_client( "prefix_padding_ms": 300, "silence_duration_ms": 500, "create_response": True, - "interrupt_response": False, + "interrupt_response": True, }, } }, @@ -233,6 +233,7 @@ async def run_client( continue if event_type == "response.done" and server_vad: + await ws.send(json.dumps({"type": "session.close"})) break if event_type == "error": diff --git a/examples/online_serving/qwen3_omni/realtime_web/__init__.py b/examples/online_serving/qwen3_omni/realtime_web/__init__.py new file mode 100644 index 000000000..14c5daad2 --- /dev/null +++ b/examples/online_serving/qwen3_omni/realtime_web/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Qwen3-Omni turn-based browser call UI.""" diff --git a/examples/online_serving/qwen3_omni/realtime_web/__main__.py b/examples/online_serving/qwen3_omni/realtime_web/__main__.py new file mode 100644 index 000000000..698179528 --- /dev/null +++ b/examples/online_serving/qwen3_omni/realtime_web/__main__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from examples.online_serving.realtime_web.server import main + +if __name__ == "__main__": + main(default_profile="qwen3-turn") diff --git a/examples/online_serving/realtime_web/README.md b/examples/online_serving/realtime_web/README.md new file mode 100644 index 000000000..854aa4718 --- /dev/null +++ b/examples/online_serving/realtime_web/README.md @@ -0,0 +1,221 @@ +# Shared Realtime Web UI + +One browser shell provides microphone capture, speaker playback, a level meter, +conversation text, and an event log. Model profiles own the wire protocol. +This implements [RFC #7222](https://github.com/vllm-project/vllm-omni/issues/7222). + +| Profile | Turn control | Camera | Playback acknowledgement | +| --- | --- | --- | --- | +| `minicpm-native` | Model-controlled listen/speak, continuous audio input | Frames accompany audio | Yes | +| `qwen3-turn --stt` (default) | User presses **Send turn** | No | No | +| `qwen3-turn --vad` | Server detects trailing silence; speech interrupts replies | Sampled frames with each spoken turn | Yes | + +Qwen3 VAD uses an engine-owned duplex plugin: microphone upload continues +while replies stream, and speech can interrupt generation and playback. The +checkpoint still generates committed turns; this is not native streaming-input +KV decoding. STT remains explicit-turn and audio-only. + +## MiniCPM: existing command stays valid + +Start your MiniCPM backend using the [existing deployment instructions](../minicpmo/README.md), then: + +```bash +python -m examples.online_serving.minicpmo.realtime_web \ + --ws-backend ws://127.0.0.1:8099 \ + --ref-audio /path/to/ref_minicpm_signature.wav +``` + +The compatibility wrapper serves the shared assets with `minicpm-native`. +The current native query, `extra_body`, reference voice, continuous microphone +upload, camera frames and playback acknowledgements are retained. + +## Qwen3: explicit-turn STT adapter + +Start the backend, then the UI in a second terminal: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 + +python -m examples.online_serving.qwen3_omni.realtime_web \ + --backend ws://127.0.0.1:8091 --stt --port 7863 +``` + +Open `http://localhost:7863`, start a session, speak, and press **Send turn**. +Wait for the answer to finish playing before speaking again. The UI opens a +fresh STT connection for each turn; the displayed conversation is a local log, +and **previous turns are not provided as model history**. The STT endpoint +accepts only model selection, so system-prompt controls are hidden in this mode. + +The adapter follows the shipped `qwen3_omni/openai_realtime_client.py` and +`vllm_omni/entrypoints/openai/realtime_connection.py`: + +1. Explicit `duplex=0` selects the legacy STT handler on a turn deployment. + It does not enable STT on a duplex deployment. +2. Send `{type: "session.update", model: ...}` and `commit(final=false)`. +3. Stream mono PCM16 at 16 kHz; **Send turn** flushes buffered audio, then sends `commit(final=true)`. +4. `response.output_audio.delta.audio` contains PCM; `response.output_audio.done` + terminates this adapter, including a response with no audio. +5. The local implementation emits **model output** on `transcription.*`, so the + UI displays it as assistant text. This differs from the early RFC's suggested + user-transcript mapping. The VAD adapter also supports `response.output_text.*` + and audio-transcript events without duplicating the answer. + +## Qwen3 with Server VAD (automatic turns) + +Run commands from the repository root with the vLLM-Omni environment activated. +Prepare the pinned Silero v6.2 ONNX artifact as described in the +[Qwen3 Server VAD instructions](../qwen3_omni/README.md#realtime-websocket-client-openai_realtime_clientpy). +The backend requires ONNX Runtime and a compatible local artifact; setting +`--vad` on the UI alone does not enable VAD on the backend. + +Create an overlay named `qwen3_vad.yaml`, replacing both paths with absolute paths +on the backend host: + +```yaml +base_config: /path/to/vllm-omni/vllm_omni/deploy/qwen3_omni_duplex.yaml +duplex_session: + server_vad_model_path: /path/to/silero_vad.onnx +``` + +The base configuration enables duplex mode and sets the per-prompt audio and image limits. + +Start the backend: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --deploy-config qwen3_vad.yaml +``` + +In a second terminal, wait until the health check succeeds, then start the UI: + +```bash +curl --fail http://127.0.0.1:8091/health +python -m examples.online_serving.qwen3_omni.realtime_web \ + --backend ws://127.0.0.1:8091 --vad --port 7863 +``` + +Open `http://localhost:7863` and start a session. Speak and pause: the server +commits the turn automatically after 500 ms of silence. No **Send turn** action +is needed. Defaults are a speech threshold of 0.5 and 300 ms of prefix padding. +Speak again to interrupt an answer. The engine owns input buffering, VAD, +request cancellation, playback acknowledgements, and conversation history. +Each committed utterance starts a normal Thinker → Talker → Code2Wav request; +it does not append input to a live Qwen KV cache. The capability response reports +`implementation_level: turn_based_duplex` and `supports_core_resumable_request: false`. + +The plugin retains up to four audio inputs (including the current turn), pruning +older audio at an 8 MiB base64 budget; one current utterance is bounded by the +engine's pending-input limit. The UI reports playback progress both on normal +drain and on interruption. Qwen does not provide text/audio alignment, so assistant +text enters history only after the complete audio response has been generated and +acknowledged as played. A partially played answer contributes no guessed text +prefix. Image conversation items remain available until deleted. Model context limits still apply. +The bundled duplex deploy allows four audio inputs and eight images per prompt. +Use headphones to avoid speaker audio triggering VAD interruption. Reference +voices and tool calls are not supported by this plugin. + +### Switching between VAD and manual turns + +| Mode | Backend | UI flag | Submit a turn | +| --- | --- | --- | --- | +| Without VAD | Default Qwen turn deployment | `--stt` (default) | Press **Send turn** | +| With VAD | Server VAD configuration and Silero artifact | `--vad` | Pause after speaking | + +Stop the existing UI process before starting another on the same port, then +refresh the browser and reconnect. To use STT, start a backend with the default +turn deployment and restart the UI with `--stt`. The duplex backend exposes +Realtime duplex and ChatCompletion; it does not expose the legacy STT handler. +To enable VAD on a default backend, first restart that backend with the overlay. +Make `--backend` match the actual backend port; the UI and backend use separate +ports. `Connection refused` means the target backend is unavailable: check its +startup log and `/health` before connecting. + +VAD connects with `duplex=1` and sends nested `session.audio.input` configuration +with `create_response: true` and `interrupt_response: true`. It sends playback +ACKs when speaker audio drains. `response.audio.done` / `response.output_audio.done` +only drain playback; `response.done` terminates a turn. Backpressure displays a +notice without pausing the microphone. Failed handshakes time out after 15 seconds; +Silero initialization failures display a specific error. + +After updating the backend Python code, restart the backend, refresh the UI, and +start a new session so the plugin and browser assets use the updated code. A +successful `/health` check confirms readiness, not which deployment mode is active. + +## Qwen VAD camera input + +After starting a session with `--vad`, click **Camera** and allow access. The UI +sends one JPEG per second, resized to a longest side of 448 pixels, independently +of microphone uploads (including while muted). Each frame is a user conversation +item using the [OpenAI Realtime image input](https://developers.openai.com/api/docs/guides/realtime-conversations#image-inputs) shape: + +```json +{"type":"conversation.item.create","item":{"id":"camera_1","type":"message","role":"user","content":[{"type":"input_image","image_url":"data:image/jpeg;base64,..."}]}} +``` + +Images become conversation context immediately; they do not start a response. +Speak and let VAD commit your audio, or send `response.create` after adding an +image/text item. A manual `input_audio_buffer.commit` only commits audio; follow +it with `response.create` to generate. Images persist across audio commits, +interruptions, and audio-buffer clears. Remove them with `conversation.item.delete`. +The browser retains eight camera items and deletes the oldest before adding a +ninth. Camera-off stops capture but leaves those items in the conversation. + +This deployment accepts JPEG/PNG Base64 data URLs with a total image-context +budget of eight images and 4 MiB; external image URLs are not supported. Clients +must delete old image items when reaching the limit. There is no precise +video/audio timestamp alignment or native streaming-video KV update. + +### Protocol compatibility boundary + +Qwen VAD uses `session.update`, 24 kHz PCM16 `input_audio_buffer.append`, +`conversation.item.create` with `input_image`/`input_text`, and `response.create`. +The engine resamples input to Qwen's 16 kHz processing rate. Qwen rejects the +custom `video_frames` field on an append: it answers whole turns, so it has no +audio unit for a frame track to align to. MiniCPM keeps that field for its +native camera protocol, where Stage 0 interleaves frames at unit boundaries. + +This is a supported subset, not full OpenAI Realtime API compatibility. +`duplex=1`, `overlap_policy`, `playback.ack`, and `session.close` remain vLLM-Omni +extensions. The legacy `--stt` mode is unchanged and is not the standard Realtime +conversation API. This demo does not add WebRTC, semantic VAD, or tool calling. + +## Shared host options + +The shared entry point can select either profile explicitly: + +```bash +python -m examples.online_serving.realtime_web --profile qwen3-turn \ + --backend ws://127.0.0.1:8091 --stt --port 7863 +``` + +- `--backend` / `--ws-backend`: backend WebSocket origin (not the full `/v1/realtime` URL). +- `--model`: override the wrapper's model name. +- `--public-realtime-url`: optional browser-visible WebSocket URL; otherwise the + static host proxies `/v1/realtime` on the same origin. +- `--host`, `--port`: UI bind address and port (default port 7862). +- `--ref-audio`: required for MiniCPM; rejected for Qwen3. + +Microphone and camera access require `localhost` or HTTPS. For a remote backend, +use an SSH tunnel to the UI host or an HTTPS reverse proxy with WebSocket support. +An explicit `wss://` public URL retains its scheme. + +## Validation + +```bash +node --test tests/examples/test_realtime_web_profiles.cjs +pytest tests/examples/test_minicpmo_realtime_web_server.py \ + tests/examples/test_minicpmo_realtime_web_static.py \ + tests/examples/test_realtime_web_profiles.py +pytest -m cpu tests/model_executor/models/qwen3_omni/test_duplex_plugin.py \ + tests/engine/duplex/test_session_runner.py \ + tests/engine/test_duplex_orchestrator.py +``` + +These CPU tests cover profile messages, two STT turns, VAD terminal/drain ordering, +backpressure, native capture during playback, server config injection, and shared +worklets. Backend regressions cover current audio inclusion in history-based +prompts, delayed playback ACK ordering, and no extra generation during silence +after a completed VAD turn. Hardware acceptance still requires a MiniCPM call +with camera/barge-in and two Qwen turns against each enabled backend, including +speech interruption in VAD mode; simulated tests do not +establish model quality or hardware latency. diff --git a/examples/online_serving/realtime_web/__init__.py b/examples/online_serving/realtime_web/__init__.py new file mode 100644 index 000000000..9e82f8967 --- /dev/null +++ b/examples/online_serving/realtime_web/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared browser audio UI and model-specific realtime profiles.""" diff --git a/examples/online_serving/realtime_web/__main__.py b/examples/online_serving/realtime_web/__main__.py new file mode 100644 index 000000000..618f3c542 --- /dev/null +++ b/examples/online_serving/realtime_web/__main__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from .server import main + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/minicpmo/realtime_web/app/index.html b/examples/online_serving/realtime_web/app/index.html similarity index 76% rename from examples/online_serving/minicpmo/realtime_web/app/index.html rename to examples/online_serving/realtime_web/app/index.html index 46ab15907..a38a731e4 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/index.html +++ b/examples/online_serving/realtime_web/app/index.html @@ -4,37 +4,35 @@ - MiniCPM-o 4.5 Voice + Omni Realtime Voice

-

Experimental full-duplex runtime

-

MiniCPM-o 4.5 Voice

+

+

Omni Realtime Voice

Offline
+

+
-
+
- +
Mic @@ -53,7 +51,7 @@

Live transcript

Conversation

- Model-controlled listen / speak +

Start a session and speak naturally.

@@ -73,7 +71,10 @@
- + + + + diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/realtime_web/app/static/app.js similarity index 59% rename from examples/online_serving/minicpmo/realtime_web/app/static/app.js rename to examples/online_serving/realtime_web/app/static/app.js index 4d2b78247..4d9215ba0 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/app.js +++ b/examples/online_serving/realtime_web/app/static/app.js @@ -1,8 +1,10 @@ (() => { 'use strict'; - const config = window.FULL_DUPLEX_CONFIG || {}; + const config = window.OMNI_REALTIME_CONFIG || window.FULL_DUPLEX_CONFIG || {}; + const profile = window.OmniRealtimeProfiles[config.profile || 'minicpm-native'](config); const callButton = document.getElementById('callButton'); + const sendTurnButton = document.getElementById('sendTurnButton'); const muteButton = document.getElementById('muteButton'); const cameraButton = document.getElementById('cameraButton'); const cameraPreview = document.getElementById('cameraPreview'); @@ -27,20 +29,26 @@ const INITIAL_PLAYBACK_BUFFER_MS = 400; const SESSION_CLOSE_TIMEOUT_MS = 1000; - // Default prompts mirroring the official MiniCPM-o-Demo presets - // (assets/presets/{omni,audio_duplex}/*.yaml). - const PROMPT_PRESETS = { - omni: 'Streaming Omni Conversation.', - chinese_call: '扮演一个具有以上声音特征的助手。请认真、高质量地回复用户的问题。' - + '请用高自然度的方式和用户聊天。你处于双工模式，可以一边听、一边说。' - + '你是由面壁智能开发的人工智能助手：面壁小钢炮。', - english_call: 'Replicate the tone and style from the input audio. Your task is to be ' - + 'a helpful assistant using this voice pattern. Please answer the user\'s questions ' - + 'seriously and in a high quality. Please chat with the user in a high naturalness ' - + 'style. You are in duplex mode, where you can listen and speak at the same time.', - }; + const PROMPT_PRESETS = profile.presets; + document.title = profile.title; + document.getElementById('pageTitle').textContent = profile.title; + document.getElementById('pageEyebrow').textContent = profile.eyebrow; + document.getElementById('profileDescription').textContent = profile.description; + document.getElementById('policyLabel').textContent = profile.policy; + cameraButton.hidden = !profile.camera; + sendTurnButton.hidden = !profile.clientCommit; + promptPreset.replaceChildren(); + for (const name of [...Object.keys(PROMPT_PRESETS), 'custom']) { + const option = document.createElement('option'); + option.value = name; + option.textContent = name.replaceAll('_', ' '); + promptPreset.appendChild(option); + } + systemPromptInput.value = Object.values(PROMPT_PRESETS)[0]; + document.getElementById('promptControls').hidden = !profile.instructions; let socket = null; + let closingSocket = null; let mediaStream = null; let captureContext = null; let captureNode = null; @@ -65,6 +73,63 @@ let liveUserTurn = null; let liveAssistantTurn = null; let sessionCloseResolver = null; + let responseComplete = false; + let playbackComplete = true; + let turnSubmitted = false; + let echoTimer = null; + let sessionGeneration = 0; + let audioChain = Promise.resolve(); + let assistantTextChannel = null; + let connectionReady = false; + let turnCounter = 0; + let stopping = null; + let turnTimeout = null; + + function markBusy() { + assistantActive = true; + if (profile.halfDuplex) pendingCapture = []; + } + + function armTurnTimeout() { + clearTimeout(turnTimeout); + turnTimeout = window.setTimeout(() => { + failSession('Timed out waiting for the model response. Check the backend event log and restart the session.'); + }, 120000); + } + + async function failSession(message) { + appendLog(message, true); + await stopSession({ terminal: false }); + setConnection('Error', 'error'); + runtimeDetail.textContent = message; + } + + function finishResponseIfReady() { + if (!responseComplete || !playbackComplete || echoTimer !== null) return; + const generation = sessionGeneration; + echoTimer = window.setTimeout(async () => { + echoTimer = null; + if (!running || generation !== sessionGeneration) return; + if (profile.reconnectEachTurn) { + connectionReady = false; + const previous = socket; + socket = null; + if (previous) { previous.onclose = null; previous.onmessage = null; previous.close(); } + try { await openSocket(); } + catch (error) { await failSession(error.message); return; } + if (!running || generation !== sessionGeneration) return; + } + assistantActive = false; + responseComplete = false; + turnSubmitted = false; + currentResponseId = null; + responseHasAudio = false; + assistantTextChannel = null; + if (profile.halfDuplex) pendingCapture = []; + sendTurnButton.disabled = !profile.clientCommit; + setModel(profile.waiting); + }, ECHO_GUARD_MS); + } function staticAssetUrl(path) { const version = String(config.appVersion || '').trim(); @@ -82,12 +147,7 @@ } function realtimeUrl() { - const url = new URL(config.realtimePath, window.location.href); - url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - url.searchParams.set('duplex', '1'); - url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5'); - url.searchParams.set('autostart', '0'); - return url.toString(); + return profile.url(config, window.location.href); } function setConnection(label, kind) { @@ -250,7 +310,7 @@ } function microphoneUploadEnabled() { - return running && !muted; + return running && connectionReady && !muted && (!profile.halfDuplex || !assistantActive); } function flushCapture() { @@ -267,19 +327,9 @@ offset += chunk.length; } pendingCapture = []; - const pcm = resampleInt16(merged, captureRate, INPUT_RATE); - const appendEvent = { - type: 'input_audio_buffer.append', - audio: int16ToBase64(pcm), - format: 'pcm16', - sample_rate_hz: INPUT_RATE, - }; - // Omni duplex: ~1 fps camera frame rides the audio append (official - // MiniCPM-o-Demo contract: one base64 JPEG per ~1 s chunk). - if (cameraPendingFrame) { - appendEvent.video_frames = [cameraPendingFrame]; - cameraPendingFrame = null; - } + const pcm = resampleInt16(merged, captureRate, profile.inputSampleRate || INPUT_RATE); + const appendEvent = profile.append(int16ToBase64(pcm), cameraPendingFrame); + cameraPendingFrame = null; socket.send(JSON.stringify(appendEvent)); } @@ -287,13 +337,17 @@ currentResponseId = responseId || currentResponseId; responseHasAudio = false; assistantActive = true; - setModel('Speaking'); + responseComplete = false; + playbackComplete = true; + if (profile.halfDuplex) pendingCapture = []; + setModel(profile.halfDuplex ? 'Thinking' : 'Speaking'); } function feedPlayback(decoded, responseId) { if (!decoded || !decoded.pcm || decoded.pcm.length === 0 || !playbackNode) return; const pcm = resampleInt16(decoded.pcm, decoded.sourceRate, playbackRate); responseHasAudio = true; + playbackComplete = false; assistantActive = true; setPlayback('Buffering'); playbackNode.port.postMessage({ @@ -305,6 +359,7 @@ } function requestPlaybackDrain(responseId) { + if (!responseHasAudio) { playbackComplete = true; finishResponseIfReady(); return; } if (!playbackNode) return; playbackNode.port.postMessage({ type: 'drain', responseId: responseId || currentResponseId }); } @@ -314,90 +369,118 @@ if (!responseId && playedMs > 0) appendLog('playback ack skipped: missing response id', true); return; } - socket.send(JSON.stringify({ - type: 'playback.ack', - response_id: responseId, - item_id: `item_${responseId}`, - played_ms: playedMs, - committed_ms: playedMs, - })); + const ack = profile.ack(responseId, playedMs); + if (ack) socket.send(JSON.stringify(ack)); } function playbackDrained(message) { const responseId = message.responseId || currentResponseId; - sendPlaybackAck(responseId, Number(message.playedMs) || 0); + if (currentResponseId && responseId !== currentResponseId) return; + if (profile.playbackAck) sendPlaybackAck(responseId, Number(message.playedMs) || 0); setPlayback('Idle'); - if (message.underrunMs > 0) { - appendLog(`playback underrun ${message.underrunMs} ms`); - } - window.setTimeout(() => { - assistantActive = false; - currentResponseId = null; - responseHasAudio = false; - if (running) setModel('Listening'); - }, ECHO_GUARD_MS); + playbackComplete = true; + if (message.underrunMs > 0) appendLog(`playback underrun ${message.underrunMs} ms`); + if (!profile.halfDuplex && !profile.waitForResponseDone) responseComplete = true; + finishResponseIfReady(); } - function handleEvent(event) { + async function handleAudioEvent(action) { + markBusy(); + currentResponseId = action.responseId || currentResponseId || `turn-${turnCounter}`; + setModel('Speaking'); + const generation = sessionGeneration; + const decoded = await decodeAudioDelta(action.event); + if (generation === sessionGeneration) feedPlayback(decoded, currentResponseId); + } + + function handleTranscriptEvent(action) { + if (action.role === 'assistant') { + // Qwen may emit both text and audio-transcript representations. + // Display one channel per response instead of duplicating the answer. + if (profile.deduplicateTranscript && assistantTextChannel && action.channel !== assistantTextChannel) return; + if (action.text) assistantTextChannel = action.channel || assistantTextChannel; + } + if (action.kind === 'text') addTranscript(action.role, action.text); + else if (profile.deduplicateTranscript && action.role === 'assistant') { + if (action.text) { + const turn = ensureTurn('assistant'); + turn.value = action.text; + turn.text.textContent = action.text; + } + } else finishTranscript(action.role, action.text); + } + + async function handleEvent(event) { appendEventLog(event); - const responseId = responseIdOf(event); - switch (event.type) { - case 'session.created': - case 'session.updated': + const action = profile.mapEvent(event); + const responseId = action.responseId; + switch (action.kind) { + case 'connected': setConnection('Connected', 'online'); - setModel('Listening'); + if (!assistantActive) setModel(profile.waiting); break; - case 'response.listen': + case 'interrupt': + sessionGeneration += 1; + // A cancelled response never reaches 'done', so nothing else would + // disarm the turn watchdog it left armed. + clearTimeout(turnTimeout); + if (playbackNode) playbackNode.port.postMessage({ type: 'clear' }); + responseComplete = true; + playbackComplete = true; + responseHasAudio = false; assistantActive = false; - setModel('Listening'); + currentResponseId = null; + assistantTextChannel = null; + finishTranscript('assistant'); + setPlayback('Idle'); + setModel(profile.waiting); break; - case 'response.created': - case 'response.speak': + case 'listen': + assistantActive = false; + setModel(profile.waiting); + break; + case 'begin': + if (profile.halfDuplex) armTurnTimeout(); + if (echoTimer !== null) { clearTimeout(echoTimer); echoTimer = null; } beginAssistant(responseId); break; - case 'response.output_audio.delta': - currentResponseId = responseId || currentResponseId; - assistantActive = true; - setModel('Speaking'); - decodeAudioDelta(event) - .then((decoded) => feedPlayback(decoded, responseId)) - .catch((error) => appendLog(`audio decode failed: ${error.message || error}`, true)); + case 'audio': + await handleAudioEvent(action); break; - case 'response.output_audio.done': + case 'drain': requestPlaybackDrain(responseId); break; - case 'response.output_audio_transcript.delta': - addTranscript('assistant', event.delta || ''); + case 'text': case 'text-final': + handleTranscriptEvent(action); break; - case 'response.output_audio_transcript.done': - finishTranscript('assistant', event.transcript || ''); - break; - case 'conversation.item.input_audio_transcription.delta': - addTranscript('user', event.delta || ''); - break; - case 'conversation.item.input_audio_transcription.completed': - finishTranscript('user', event.transcript || ''); - break; - case 'response.done': + case 'done': + clearTimeout(turnTimeout); + responseComplete = true; finishTranscript('assistant'); - if (!responseHasAudio) requestPlaybackDrain(responseId); + requestPlaybackDrain(responseId); + finishResponseIfReady(); break; - case 'playback.acknowledged': - { - const acknowledgement = event.event || event; - runtimeDetail.textContent = `Playback committed ${acknowledgement.committed_ms || 0} ms`; - } + case 'backpressure': + armTurnTimeout(); + markBusy(); + setModel('Thinking / Speaking'); + runtimeDetail.textContent = action.message; break; - case 'session.closed': + case 'ack': + runtimeDetail.textContent = `Playback committed ${action.committedMs} ms`; + break; + case 'closed': if (sessionCloseResolver) sessionCloseResolver(); sessionCloseResolver = null; break; case 'error': - setConnection('Error', 'error'); - runtimeDetail.textContent = String(event.error || event.code || 'Server error'); - break; - default: + if (action.fatal) { await failSession(action.message); break; } + // A rejected frame or a refused update costs one operation, not the + // call: report it and keep listening. + appendLog(action.message, true); + runtimeDetail.textContent = action.message; break; + default: break; } } @@ -406,8 +489,15 @@ playbackRate = playbackContext.sampleRate; await playbackContext.audioWorklet.addModule(staticAssetUrl('static/playback_worklet.js')); playbackNode = new AudioWorkletNode(playbackContext, 'fullduplex-pcm-playback'); + const currentPlaybackNode = playbackNode; playbackNode.port.onmessage = (message) => { + if (playbackNode !== currentPlaybackNode) return; if (message.data.type === 'playback-started') setPlayback('Playing'); + else if (message.data.type === 'playback-stopped') { + // The old response may report after the next one has started. ACK its + // own cursor without changing the current response's UI state. + if (profile.playbackAck) sendPlaybackAck(message.data.responseId, Number(message.data.playedMs) || 0); + } else if (message.data.type === 'playback-drained') playbackDrained(message.data); else if (message.data.type === 'playback-underrun') { runtimeDetail.textContent = `Playback underrun ${message.data.underrunMs || 0} ms`; @@ -451,48 +541,69 @@ function openSocket() { return new Promise((resolve, reject) => { const url = realtimeUrl(); - socket = new WebSocket(url); + const current = new WebSocket(url); + socket = current; + connectionReady = false; + turnCounter += 1; let settled = false; - socket.onopen = () => { + const timer = window.setTimeout(() => { + if (settled) return; settled = true; - const extraBody = { - auto_response: true, - }; - const session = { - modalities: ['audio', 'text'], - voice: 'default', - extra_body: extraBody, - }; - // Reference voice for TTS cloning, provided by the server via - // --ref-audio (mirrors the official demo's default ref audio). - if (config.refAudio) session.ref_audio = config.refAudio; - const instructions = systemPromptInput ? systemPromptInput.value.trim() : ''; - if (instructions) session.instructions = instructions; - socket.send(JSON.stringify({ type: 'session.update', session })); - runtimeDetail.textContent = `${captureRate} Hz capture / ${playbackRate} Hz playback`; + current.onclose = null; + current.close(); + reject(new Error(`Session handshake timed out. ${profile.connectionHint}`)); + }, 15000); + const rejectOnce = (message) => { + clearTimeout(timer); + if (!settled) { settled = true; reject(new Error(message)); } + }; + current.onopen = () => { + const instructions = systemPromptInput.value.trim(); + for (const event of profile.initialMessages(config, instructions)) current.send(JSON.stringify(event)); appendLog(`websocket open ${url}`); - resolve(); }; - socket.onmessage = (message) => { - if (typeof message.data !== 'string') return; - try { - handleEvent(JSON.parse(message.data)); - } catch (error) { - appendLog(`invalid server event: ${error.message || error}`, true); + current.onmessage = (message) => { + if (typeof message.data !== 'string' || (socket !== current && closingSocket !== current)) return; + let event; + try { event = JSON.parse(message.data); } + catch (error) { + if (socket !== current) return; + const detail = `Invalid server event: ${error.message}`; + if (!settled) rejectOnce(detail); + else void failSession(detail); + return; } - }; - socket.onerror = () => { - if (!settled) { + // Shutdown has detached the active socket. Only its close acknowledgement + // may pass, directly: the event queue can itself be awaiting cleanup. + if (closingSocket === current) { + if (event?.type === 'session.closed' && sessionCloseResolver) sessionCloseResolver(); + return; + } + const action = profile.mapEvent(event); + if (action.kind === 'error' && !settled) { + rejectOnce(action.message); + return; + } + if (event.type === profile.readyEvent && !settled) { + clearTimeout(timer); settled = true; - reject(new Error(`WebSocket connection failed: ${url}`)); + connectionReady = true; + runtimeDetail.textContent = `${captureRate} Hz capture / ${playbackRate} Hz playback`; + resolve(); } + // Serialize decoding with terminal events: a drain must never overtake + // an asynchronously decoded WAV chunk. + audioChain = audioChain.then(() => { + if (socket === current || action.kind === 'closed') return handleEvent(event); + }).catch((error) => failSession(`Server event failed: ${error.message}`)); }; - socket.onclose = (event) => { + current.onerror = () => rejectOnce(`WebSocket connection failed. ${profile.connectionHint}`); + current.onclose = (event) => { + clearTimeout(timer); + connectionReady = false; appendLog(`websocket closed code=${event.code}`); - if (running) { - setConnection('Disconnected', 'error'); - stopSession(); - } + if (!settled) rejectOnce(`Connection closed before session ready. ${profile.connectionHint}`); + else if (running && socket === current) failSession(`Backend disconnected (${event.code}). ${profile.connectionHint}`); }; }); } @@ -515,46 +626,59 @@ setConnection('Connecting', 'connecting'); runtimeDetail.textContent = 'Requesting microphone access'; try { + sessionGeneration += 1; + audioChain = Promise.resolve(); await openPlayback(); await openCapture(); await openSocket(); running = true; muted = false; assistantActive = false; - sendTimer = window.setInterval(flushCapture, SEND_INTERVAL_MS); + sendTimer = window.setInterval(flushCapture, profile.sendIntervalMs || SEND_INTERVAL_MS); startClock(); callButton.textContent = 'End session'; callButton.classList.add('is-active'); muteButton.disabled = false; - cameraButton.disabled = false; + cameraButton.disabled = !profile.camera; + sendTurnButton.disabled = !profile.clientCommit; setConnection('Connected', 'online'); - setModel('Listening'); + setModel(profile.waiting); appendLog('session started'); } catch (error) { appendLog(`start failed: ${error.message || error}`, true); + await stopSession({ terminal: false }); setConnection('Error', 'error'); runtimeDetail.textContent = String(error.message || error); - await stopSession(); } finally { callButton.disabled = false; } } async function startCamera() { - if (cameraStream) return; + if (!profile.camera || cameraStream) return; cameraStream = await navigator.mediaDevices.getUserMedia({ video: true, audio: false }); cameraPreview.srcObject = cameraStream; cameraPreview.style.display = ''; await cameraPreview.play().catch(() => {}); // Official omni-duplex cadence: one JPEG (quality 0.7) per ~1 s chunk, // no client-side resize (the server normalizes at scale_resolution=448). - cameraTimer = window.setInterval(() => { + const captureCameraFrame = () => { if (!cameraStream || cameraPreview.videoWidth === 0) return; - cameraCanvas.width = cameraPreview.videoWidth; - cameraCanvas.height = cameraPreview.videoHeight; - cameraCanvas.getContext('2d').drawImage(cameraPreview, 0, 0); - cameraPendingFrame = cameraCanvas.toDataURL('image/jpeg', 0.7).split(',')[1]; - }, 1000); + const scale = profile.cameraMaxDimension + ? Math.min(1, profile.cameraMaxDimension / Math.max(cameraPreview.videoWidth, cameraPreview.videoHeight)) : 1; + cameraCanvas.width = Math.max(1, Math.round(cameraPreview.videoWidth * scale)); + cameraCanvas.height = Math.max(1, Math.round(cameraPreview.videoHeight * scale)); + cameraCanvas.getContext('2d').drawImage(cameraPreview, 0, 0, cameraCanvas.width, cameraCanvas.height); + const frame = cameraCanvas.toDataURL('image/jpeg', 0.7).split(',')[1]; + if (profile.imageMessages) { + if (connectionReady && socket?.readyState === WebSocket.OPEN) { + for (const event of profile.imageMessages(frame)) socket.send(JSON.stringify(event)); + } + } else cameraPendingFrame = frame; + }; + // Do not make the first spoken turn race a one-second timer. + captureCameraFrame(); + cameraTimer = window.setInterval(captureCameraFrame, 1000); cameraButton.textContent = 'Camera off'; cameraButton.classList.add('is-active'); appendLog('camera on (1 fps omni frames)'); @@ -590,15 +714,32 @@ const finish = () => { if (done) return; done = true; + clearTimeout(timer); if (sessionCloseResolver === finish) sessionCloseResolver = null; resolve(); }; sessionCloseResolver = finish; - window.setTimeout(finish, timeoutMs); + const timer = window.setTimeout(finish, timeoutMs); }); } - async function stopSession({ terminal = true } = {}) { + function stopSession(options = {}) { + if (stopping) return stopping; + stopping = cleanupSession(options).finally(() => { stopping = null; }); + return stopping; + } + + async function cleanupSession({ terminal = true } = {}) { + sessionGeneration += 1; + connectionReady = false; + clearTimeout(echoTimer); + clearTimeout(turnTimeout); + echoTimer = null; + responseComplete = false; + playbackComplete = true; + turnSubmitted = false; + assistantTextChannel = null; + sendTurnButton.disabled = true; running = false; assistantActive = false; pendingCapture = []; @@ -607,15 +748,16 @@ sendTimer = null; clockTimer = null; if (socket) { - const closingSocket = socket; + closingSocket = socket; socket = null; closingSocket.onclose = null; - if (terminal && closingSocket.readyState === WebSocket.OPEN) { + if (terminal && profile.closeSession && closingSocket.readyState === WebSocket.OPEN) { const closed = waitForSessionClosed(closingSocket, SESSION_CLOSE_TIMEOUT_MS); closingSocket.send(JSON.stringify({ type: 'session.close' })); await closed; } closingSocket.close(1000, 'client stop'); + closingSocket = null; } if (playbackNode) playbackNode.port.postMessage({ type: 'clear' }); if (mediaStream) { @@ -632,6 +774,8 @@ playbackNode = null; currentResponseId = null; responseHasAudio = false; + finishTranscript('user'); + finishTranscript('assistant'); meterFill.style.width = '0%'; sessionTimer.textContent = '00:00'; callButton.textContent = 'Start session'; @@ -659,6 +803,18 @@ if (running) stopSession(); else startSession(); }); + sendTurnButton.addEventListener('click', () => { + if (!running || !connectionReady || !profile.clientCommit || turnSubmitted) return; + flushCapture(); + for (const event of profile.commitMessages()) socket.send(JSON.stringify(event)); + turnSubmitted = true; + responseComplete = false; + markBusy(); + armTurnTimeout(); + sendTurnButton.disabled = true; + setModel('Thinking'); + appendLog('turn submitted'); + }); muteButton.addEventListener('click', toggleMute); clearLogButton.addEventListener('click', () => { eventLog.textContent = ''; diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/pcm_worklet.js b/examples/online_serving/realtime_web/app/static/pcm_worklet.js similarity index 100% rename from examples/online_serving/minicpmo/realtime_web/app/static/pcm_worklet.js rename to examples/online_serving/realtime_web/app/static/pcm_worklet.js diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js b/examples/online_serving/realtime_web/app/static/playback_worklet.js similarity index 95% rename from examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js rename to examples/online_serving/realtime_web/app/static/playback_worklet.js index f0d3d42ae..8aaa131c6 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js +++ b/examples/online_serving/realtime_web/app/static/playback_worklet.js @@ -39,6 +39,13 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { } this.notifyIfDrained(); } else if (message.type === 'clear') { + if (this.activeResponseId) { + this.port.postMessage({ + type: 'playback-stopped', + responseId: this.activeResponseId, + playedMs: Math.round((this.playedFrames * 1000) / sampleRate), + }); + } this.queue = []; this.offset = 0; this.playedFrames = 0; diff --git a/examples/online_serving/realtime_web/app/static/profiles/common.js b/examples/online_serving/realtime_web/app/static/profiles/common.js new file mode 100644 index 000000000..feefbaa74 --- /dev/null +++ b/examples/online_serving/realtime_web/app/static/profiles/common.js @@ -0,0 +1,56 @@ +/* Shared wire helpers. Profiles translate server events into shell actions. */ +(function (root) { + 'use strict'; + const profiles = root.OmniRealtimeProfiles = root.OmniRealtimeProfiles || {}; + profiles.url = (config, location) => { + const url = new URL(config.realtimePath || 'v1/realtime', location); + if (url.protocol === 'http:') url.protocol = 'ws:'; + if (url.protocol === 'https:') url.protocol = 'wss:'; + return url; + }; + /* Errors that leave the session unusable, so the shell has to tear it down. + Everything else costs one operation --- a rejected image, a refused + session.update, a late playback ack --- and the call carries on. + 'runtime_append_failed' is deliberately absent: the engine reports it and + keeps the session, unlike 'runtime_append_task_failed' which kills it. */ + profiles.fatalErrorCodes = new Set([ + 'unknown_session', + 'session_closed', + 'resource_exhausted', + 'unsupported', + 'unsupported_turn_detection', + 'server_vad_unavailable', + 'runtime_append_task_failed', + 'internal_error', + ]); + profiles.event = (event) => { + const responseId = event.response_id || event.response?.id || null; + const action = (kind, extra = {}) => ({ kind, responseId, ...extra }); + switch (event.type) { + case 'session.created': case 'session.updated': return action('connected'); + case 'response.created': return action('begin'); + case 'response.audio.delta': case 'response.output_audio.delta': + return action('audio', { event: { ...event, delta: event.delta || event.audio || event.response?.audio } }); + case 'response.audio.done': case 'response.output_audio.done': return action('drain'); + case 'response.audio_transcript.delta': case 'response.output_audio_transcript.delta': + return action('text', { role: 'assistant', channel: 'audio', text: event.delta || '' }); + case 'response.audio_transcript.done': case 'response.output_audio_transcript.done': + return action('text-final', { role: 'assistant', channel: 'audio', text: event.transcript || '' }); + case 'conversation.item.input_audio_transcription.delta': + return action('text', { role: 'user', text: event.delta || '' }); + case 'conversation.item.input_audio_transcription.completed': + return action('text-final', { role: 'user', text: event.transcript || '' }); + case 'response.done': return action('done'); + case 'session.closed': return action('closed'); + case 'error': { + const code = event.code || event.error?.code; + return action('error', { + code, + fatal: profiles.fatalErrorCodes.has(code), + message: typeof event.error === 'string' ? event.error : event.error?.message || event.message || code || 'Server error', + }); + } + default: return action('ignore'); + } + }; +})(globalThis); diff --git a/examples/online_serving/realtime_web/app/static/profiles/minicpm_native.js b/examples/online_serving/realtime_web/app/static/profiles/minicpm_native.js new file mode 100644 index 000000000..801d549be --- /dev/null +++ b/examples/online_serving/realtime_web/app/static/profiles/minicpm_native.js @@ -0,0 +1,55 @@ +(function (root) { + 'use strict'; + const profiles = root.OmniRealtimeProfiles; + profiles['minicpm-native'] = () => ({ + title: 'MiniCPM-o 4.5 Voice', eyebrow: 'Experimental full-duplex runtime', + policy: 'Model-controlled listen / speak', + description: 'Speak naturally. The model can listen and speak at the same time. Camera frames are sent with your audio.', + waiting: 'Listening', camera: true, playbackAck: true, clientCommit: false, + halfDuplex: false, closeSession: true, reconnectEachTurn: false, + readyEvent: 'session.updated', instructions: true, sendIntervalMs: 200, + presets: { + omni: 'Streaming Omni Conversation.', + chinese_call: '扮演一个具有以上声音特征的助手。请认真、高质量地回复用户的问题。' + + '请用高自然度的方式和用户聊天。你处于双工模式，可以一边听、一边说。' + + '你是由面壁智能开发的人工智能助手：面壁小钢炮。', + english_call: 'Replicate the tone and style from the input audio. Your task is to be ' + + 'a helpful assistant using this voice pattern. Please answer the user\'s questions ' + + 'seriously and in a high quality. Please chat with the user in a high naturalness ' + + 'style. You are in duplex mode, where you can listen and speak at the same time.', + }, + url(config, location) { + const url = profiles.url(config, location); + url.searchParams.set('duplex', '1'); + url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5'); + url.searchParams.set('autostart', '0'); + return url.toString(); + }, + initialMessages(config, instructions) { + const session = { modalities: ['audio', 'text'], voice: 'default', + extra_body: { auto_response: true } }; + if (config.refAudio) session.ref_audio = config.refAudio; + if (instructions) session.instructions = instructions; + return [{ type: 'session.update', session }]; + }, + append(audio, frame) { + const event = { type: 'input_audio_buffer.append', audio, format: 'pcm16', sample_rate_hz: 16000 }; + if (frame) event.video_frames = [frame]; + return event; + }, + commitMessages: () => [], + ack(responseId, playedMs) { + return { type: 'playback.ack', response_id: responseId, item_id: `item_${responseId}`, + played_ms: playedMs, committed_ms: playedMs }; + }, + mapEvent(event) { + if (event.type === 'response.listen') return { kind: 'listen' }; + if (event.type === 'response.speak') return { kind: 'begin', responseId: event.response_id || event.response?.id }; + if (event.type === 'playback.acknowledged') { + return { kind: 'ack', committedMs: (event.event || event).committed_ms || 0 }; + } + return profiles.event(event); + }, + connectionHint: 'Check the MiniCPM duplex deployment and reference voice configuration.', + }); +})(globalThis); diff --git a/examples/online_serving/realtime_web/app/static/profiles/qwen3_turn.js b/examples/online_serving/realtime_web/app/static/profiles/qwen3_turn.js new file mode 100644 index 000000000..d6df32edd --- /dev/null +++ b/examples/online_serving/realtime_web/app/static/profiles/qwen3_turn.js @@ -0,0 +1,100 @@ +(function (root) { + 'use strict'; + const profiles = root.OmniRealtimeProfiles; + profiles['qwen3-turn'] = (config) => { + const vad = config.adapter === 'vad'; + let imageSeq = 0; + let imageItems = []; + // The engine refuses a ninth image (or 4 MiB of them) for the whole + // session, so the camera track retires its own oldest frame instead. + const MAX_IMAGE_ITEMS = 8; + const hint = 'This Qwen deploy needs session_mode: duplex and Server VAD support; use --stt or enable Server VAD.'; + return { + title: 'Qwen3-Omni Voice', eyebrow: vad ? 'Engine-owned duplex call' : 'Turn-based realtime call', + policy: vad ? 'Server VAD • interruptible responses' : 'Send turn • one response at a time', + description: vad + ? 'Pause to send your turn; speak again to interrupt. Microphone upload continues during playback. Requires the Qwen duplex plugin and Silero. Optional camera frames accompany each spoken turn.' + : 'Speak, then press Send turn. Microphone upload pauses while the model answers. Each turn uses a fresh connection; conversation history is not carried between turns. No barge-in or camera input.', + waiting: 'Waiting for you', camera: vad, cameraMaxDimension: 448, playbackAck: vad, clientCommit: !vad, + waitForResponseDone: true, halfDuplex: !vad, deduplicateTranscript: true, closeSession: vad, reconnectEachTurn: !vad, + inputSampleRate: vad ? 24000 : 16000, + readyEvent: vad ? 'session.updated' : 'session.created', instructions: vad, sendIntervalMs: 200, + presets: { assistant: 'You are a helpful assistant. Answer clearly and concisely.' }, + url(config, location) { + const url = profiles.url(config, location); + for (const key of ['native_duplex', 'minicpmo45_native_duplex', 'autostart']) url.searchParams.delete(key); + // Explicit opt-out matches the local file client, including servers + // that auto-select the duplex endpoint when the query is omitted. + url.searchParams.set('duplex', vad ? '1' : '0'); + url.searchParams.set('model', config.model || 'Qwen/Qwen3-Omni-30B-A3B-Instruct'); + return url.toString(); + }, + initialMessages(config, instructions) { + imageItems = []; + if (!vad) return [ + { type: 'session.update', model: config.model }, + { type: 'input_audio_buffer.commit', final: false }, + ]; + const session = { model: config.model, overlap_policy: 'barge_in_on_speech', audio: { input: { + format: { type: 'audio/pcm', rate: 24000 }, + turn_detection: { type: 'server_vad', threshold: 0.5, prefix_padding_ms: 300, + silence_duration_ms: 500, create_response: true, interrupt_response: true }, + } } }; + if (instructions) session.instructions = instructions; + return [{ type: 'session.update', session }]; + }, + append(audio) { + return vad ? { type: 'input_audio_buffer.append', audio } + : { type: 'input_audio_buffer.append', audio, format: 'pcm16', sample_rate_hz: 16000 }; + }, + // Qwen is a turn model: a frame is a picture added to the conversation, + // not a track interleaved into the audio. That is exactly the OpenAI + // Realtime image interface, so the camera sends conversation items + // rather than an extension field on the append. + imageMessages(frame) { + if (!vad) return []; + const messages = []; + if (imageItems.length >= MAX_IMAGE_ITEMS) { + // Retire before creating: the engine counts what is already stored. + messages.push({ type: 'conversation.item.delete', item_id: imageItems.shift() }); + } + const id = `camera_${++imageSeq}`; + imageItems.push(id); + messages.push({ type: 'conversation.item.create', item: { + id, type: 'message', role: 'user', + content: [{ type: 'input_image', image_url: `data:image/jpeg;base64,${frame}` }], + } }); + return messages; + }, + commitMessages: () => vad ? [] : [{ type: 'input_audio_buffer.commit', final: true }], + ack: (responseId, playedMs) => vad ? { type: 'playback.ack', response_id: responseId, + item_id: `item_${responseId}`, played_ms: playedMs, committed_ms: playedMs } : null, + mapEvent(event) { + if (event.type === 'playback.acknowledged') return { kind: 'ack', committedMs: (event.event || event).committed_ms || 0 }; + if (event.type === 'response.output_text.delta' || event.type === 'transcription.delta') { + // On the shipped STT route transcription.* is model-generated text, + // not a separate ASR transcript of the user. See realtime_connection.py. + return { kind: 'text', role: 'assistant', channel: 'text', text: event.delta || '' }; + } + if (event.type === 'response.output_text.done' || event.type === 'transcription.done') { + return { kind: 'text-final', role: 'assistant', channel: 'text', text: event.text || '' }; + } + if (vad && (event.type === 'output_audio_buffer.cleared' || + (event.type === 'response.done' && event.response?.status === 'cancelled'))) { + return { kind: 'interrupt' }; + } + if (event.type === 'input_audio_buffer.committed') return { kind: 'begin' }; + if (event.type === 'input_audio_buffer.cleared') return { kind: 'backpressure', message: 'Input cleared. Model is answering; please wait.' }; + if (!vad && event.type === 'response.output_audio.done') return { kind: 'done' }; + const action = profiles.event(event); + if (action.kind === 'error') { + if (action.code === 'input_backpressure') return { ...action, kind: 'backpressure', message: 'Input buffer is full; please repeat dropped speech after the response.' }; + if (action.code === 'server_vad_initialization_failed') action.message = 'Server VAD could not load Silero. Configure the server artifact or restart this UI with --stt.'; + else if (vad && action.code === 'unsupported') action.message += ` ${hint}`; + } + return action; + }, + connectionHint: vad ? hint : 'Check the Qwen /v1/realtime STT endpoint and model name.', + }; + }; +})(globalThis); diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/styles.css b/examples/online_serving/realtime_web/app/static/styles.css similarity index 98% rename from examples/online_serving/minicpmo/realtime_web/app/static/styles.css rename to examples/online_serving/realtime_web/app/static/styles.css index 9b080176b..bf25c3eff 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/styles.css +++ b/examples/online_serving/realtime_web/app/static/styles.css @@ -125,3 +125,6 @@ button:disabled { opacity: .46; cursor: default; } @media (prefers-reduced-motion: reduce) { *, *::before, *::after { animation-duration: .01ms !important; animation-iteration-count: 1 !important; } } + +/* Profile-specific controls may have inline flex styles. */ +[hidden] { display: none !important; } diff --git a/examples/online_serving/realtime_web/server.py b/examples/online_serving/realtime_web/server.py new file mode 100644 index 000000000..85e4b63d5 --- /dev/null +++ b/examples/online_serving/realtime_web/server.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Static host and same-origin proxy for the shared realtime browser UI.""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import contextlib +import hashlib +import json +import logging +from pathlib import Path +from urllib.parse import urlencode + +import uvicorn +import websockets +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.responses import HTMLResponse, Response +from fastapi.staticfiles import StaticFiles + +logger = logging.getLogger(__name__) +APP_DIR = Path(__file__).parent / "app" +STATIC_DIR = APP_DIR / "static" + + +def _join_ws_url(base: str, path: str, query: str) -> str: + return base.rstrip("/") + path + (("?" + query) if query else "") + + +async def _pump_client_to_backend(client: WebSocket, backend) -> None: + try: + while True: + message = await client.receive() + if message["type"] == "websocket.disconnect": + await backend.close() + return + if message.get("text") is not None: + await backend.send(message["text"]) + elif message.get("bytes") is not None: + await backend.send(message["bytes"]) + except (WebSocketDisconnect, websockets.ConnectionClosed): + with contextlib.suppress(Exception): + await backend.close() + + +async def _pump_backend_to_client(client: WebSocket, backend) -> None: + try: + async for message in backend: + if isinstance(message, bytes): + await client.send_bytes(message) + else: + await client.send_text(message) + except (WebSocketDisconnect, websockets.ConnectionClosed): + return + except RuntimeError as exc: + if "websocket.send" in str(exc) and "websocket.close" in str(exc): + return + raise + + +def _expected_proxy_close(exc: BaseException) -> bool: + if isinstance(exc, (WebSocketDisconnect, websockets.ConnectionClosed, asyncio.CancelledError)): + return True + return isinstance(exc, RuntimeError) and "websocket.send" in str(exc) and "websocket.close" in str(exc) + + +def build_app( + *, + ws_backend: str | None = None, + model: str | None = None, + public_realtime_url: str | None = None, + ref_audio: str | None = None, + profile: str = "minicpm-native", + adapter: str = "stt", +) -> FastAPI: + if profile not in {"minicpm-native", "qwen3-turn"}: + raise ValueError(f"Unknown profile: {profile}") + if adapter not in {"stt", "vad"}: + raise ValueError(f"Unknown adapter: {adapter}") + if profile == "qwen3-turn" and ref_audio: + raise ValueError("Qwen3 profile does not accept --ref-audio") + native = profile == "minicpm-native" + ws_backend = ws_backend or ("ws://127.0.0.1:8099" if native else "ws://127.0.0.1:8091") + model = model or ("openbmb/MiniCPM-o-4_5" if native else "Qwen/Qwen3-Omni-30B-A3B-Instruct") + app = FastAPI(title="Omni Realtime Web UI") + index_path = APP_DIR / "index.html" + app_version_hash = hashlib.sha256() + for asset_path in ( + STATIC_DIR / "app.js", + STATIC_DIR / "pcm_worklet.js", + STATIC_DIR / "playback_worklet.js", + *sorted((STATIC_DIR / "profiles").glob("*.js")), + ): + app_version_hash.update(asset_path.read_bytes()) + app_version = app_version_hash.hexdigest()[:12] + + ref_audio_uri: str | None = None + if ref_audio: + ref_path = Path(ref_audio) + if not ref_path.is_file(): + raise SystemExit(f"--ref-audio not found: {ref_audio}") + encoded = base64.b64encode(ref_path.read_bytes()).decode("ascii") + ref_audio_uri = f"data:audio/wav;base64,{encoded}" + + @app.get("/", response_class=HTMLResponse) + def index() -> HTMLResponse: + config = json.dumps( + { + "model": model, + "profile": profile, + "adapter": adapter, + "realtimePath": public_realtime_url or "v1/realtime", + "refAudio": ref_audio_uri, + "appVersion": app_version, + }, + ensure_ascii=True, + ).replace("<", "\\u003c") + html = ( + index_path.read_text(encoding="utf-8") + .replace("__FULL_DUPLEX_CONFIG__", config) + .replace("__FULL_DUPLEX_APP_VERSION__", app_version) + ) + return HTMLResponse(html, headers={"Cache-Control": "no-store"}) + + @app.get("/healthz") + def healthz() -> Response: + return Response(content="ok", media_type="text/plain") + + @app.websocket("/v1/realtime") + async def realtime_proxy(websocket: WebSocket) -> None: + await websocket.accept() + query = urlencode(websocket.query_params.multi_items()) + backend_url = _join_ws_url(ws_backend, "/v1/realtime", query) + logger.info("Proxying Realtime WebSocket to %s", backend_url) + try: + async with websockets.connect( + backend_url, + max_size=64 * 1024 * 1024, + ) as backend: + tasks = { + asyncio.create_task(_pump_client_to_backend(websocket, backend)), + asyncio.create_task(_pump_backend_to_client(websocket, backend)), + } + done, pending = await asyncio.wait( + tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + for result in await asyncio.gather(*done, return_exceptions=True): + if isinstance(result, BaseException) and not _expected_proxy_close(result): + raise result + except (WebSocketDisconnect, websockets.ConnectionClosed): + return + except (OSError, TimeoutError) as exc: + message = ( + f"Cannot connect to realtime backend {ws_backend}. " + "Check that vllm serve is running and --backend matches its --port. " + "For containers or remote hosts, use the backend address reachable from the UI server." + ) + logger.warning("%s (%s)", message, exc) + with contextlib.suppress(Exception): + await websocket.send_json({"type": "error", "code": "backend_unavailable", "error": message}) + await websocket.close(code=1011) + except Exception: + logger.exception("Realtime WebSocket proxy failed") + with contextlib.suppress(Exception): + await websocket.close(code=1011) + finally: + with contextlib.suppress(Exception): + await websocket.close() + + app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") + return app + + +def main(default_profile: str = "minicpm-native") -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=7862) + parser.add_argument("--profile", choices=("minicpm-native", "qwen3-turn"), default=default_profile) + parser.add_argument("--ws-backend", "--backend", dest="ws_backend") + mode = parser.add_mutually_exclusive_group() + mode.add_argument("--stt", dest="adapter", action="store_const", const="stt") + mode.add_argument("--vad", dest="adapter", action="store_const", const="vad") + parser.set_defaults(adapter="stt") + parser.add_argument( + "--public-realtime-url", + help="Browser-visible ws:// or wss:// Realtime URL; defaults to the same-origin proxy.", + ) + parser.add_argument("--model") + parser.add_argument( + "--ref-audio", + help=( + "Reference voice wav for TTS voice cloning, e.g. the " + "official MiniCPM-o-Demo assets/ref_audio/ref_minicpm_signature.wav" + ), + ) + args = parser.parse_args() + native = args.profile == "minicpm-native" + if native and not args.ref_audio: + parser.error("--ref-audio is required for the minicpm-native profile") + if not native and args.ref_audio: + parser.error("--ref-audio is only supported by minicpm-native") + args.model = args.model or ("openbmb/MiniCPM-o-4_5" if native else "Qwen/Qwen3-Omni-30B-A3B-Instruct") + args.ws_backend = args.ws_backend or ("ws://127.0.0.1:8099" if native else "ws://127.0.0.1:8091") + + logging.basicConfig(level=logging.INFO) + uvicorn.run( + build_app( + ws_backend=args.ws_backend, + model=args.model, + public_realtime_url=args.public_realtime_url, + ref_audio=args.ref_audio, + profile=args.profile, + adapter=args.adapter, + ), + host=args.host, + port=args.port, + log_level="info", + ) + + +if __name__ == "__main__": + main() diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index b73062689..eb073f0c7 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 28, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 64, + EnvironmentVariableCategory.MODEL_SPECIFIC: 65, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 15, + ModelEnvironmentVariableDisposition.INTERNALIZE: 16, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/engine/duplex/test_engine_session.py b/tests/engine/duplex/test_engine_session.py index fccaa59e8..c90291a81 100644 --- a/tests/engine/duplex/test_engine_session.py +++ b/tests/engine/duplex/test_engine_session.py @@ -567,10 +567,10 @@ def test_overlap_policy_defaults_and_invalid_values_to_listen_only(): assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only" -def test_capabilities_as_dict_reports_model_native_duplex_constants(): +def test_capabilities_as_dict_reports_declared_implementation(): caps = DuplexCapabilities().as_dict() - assert caps["implementation_level"] == "model_native_duplex" + assert caps["implementation_level"] == "turn_based_duplex" assert caps["input_modes"] == ["append_audio_chunk"] assert caps["supports_kv_lease"] is False assert caps["supports_core_kv_lease"] is False diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py index 4f1962ee5..bccce18f1 100644 --- a/tests/engine/duplex/test_session_manager.py +++ b/tests/engine/duplex/test_session_manager.py @@ -196,7 +196,7 @@ class FakePlugin(DuplexModelPlugin): def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: del max_sessions - return DuplexCapabilities(supports_input_append=True) + return DuplexCapabilities(supports_input_append=True, supports_core_resumable_request=True) def validate_client_extra_body(self, extra_body: object) -> None: pass @@ -408,7 +408,9 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non assert result.control_id == "open-sid-open" assert result.session_id == "sid-open" assert result.lease_generation == 0 - assert result.capabilities == DuplexCapabilities(supports_input_append=True) + assert result.capabilities == DuplexCapabilities( + supports_input_append=True, supports_core_resumable_request=True + ) assert result.public_session is not None assert result.public_session["id"] == "sid-open" assert result.public_session["voice"] == "test" diff --git a/tests/engine/duplex/test_vad_backend.py b/tests/engine/duplex/test_vad_backend.py index 8a22259ac..f3163ff4e 100644 --- a/tests/engine/duplex/test_vad_backend.py +++ b/tests/engine/duplex/test_vad_backend.py @@ -107,8 +107,8 @@ def test_endpoint_decisions_match_the_serving_side_policy( # --------------------------------------------------------------------------- # -def test_a_loud_frame_does_not_restart_the_silence_timer() -> None: - """Silero v6.2 hysteresis: once silence is running, a loud frame only delays it. +def test_hysteresis_band_does_not_restart_the_silence_timer() -> None: + """Scores between the negative and activation thresholds preserve the timer. The frame at 0.45 sits above the negative threshold (0.35) but below the activation threshold. It must not cancel the pending endpoint, or a speaker @@ -164,9 +164,14 @@ def _pcm16_b64(samples: np.ndarray) -> str: def test_pcm16_input_is_accepted_alongside_float32() -> None: config = SileroVADConfig(threshold=0.5) - tone = np.zeros(FRAME * 2, dtype=np.float32) + tone: np.ndarray = np.zeros(FRAME * 2, dtype=np.float32) seen: list[int] = [] - vad = SileroStreamingVAD(config, frame_scorer=lambda frame: seen.append(frame.size) or 0.0) + + def _score_frame(frame: np.ndarray) -> float: + seen.append(frame.size) + return 0.0 + + vad = SileroStreamingVAD(config, frame_scorer=_score_frame) vad.process_base64(_pcm16_b64(tone), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ) assert seen == [FRAME, FRAME] @@ -272,7 +277,7 @@ def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tm intra_op_num_threads = 0 class FakeInferenceSession: - def __init__(self, path: str, *, providers: list[str], sess_options: object) -> None: + def __init__(self, path: str, *, providers: list[str], sess_options: FakeSessionOptions) -> None: self.providers = providers self.sess_options = sess_options self.calls: list[dict[str, np.ndarray]] = [] @@ -313,7 +318,7 @@ def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tm assert session.calls[0]["sr"].item() == SAMPLE_RATE_HZ state = backend.new_state() - first_frame = np.arange(backend.frame_samples, dtype=np.float32) + first_frame: np.ndarray = np.arange(backend.frame_samples, dtype=np.float32) probability, state = backend.infer(first_frame, state) first_call = session.calls[1] @@ -377,3 +382,17 @@ def test_the_documented_server_vad_fields_are_all_accepted() -> None: ) is None ) + + +@pytest.mark.parametrize("speech_probability", [0.5, 0.9]) +def test_resumed_speech_cancels_pending_silence(speech_probability): + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=96, min_speech_duration_ms=32) + # 32 ms pause, resumed speech, then a full 96 ms pause: one utterance. + scores = [0.9, 0.1, speech_probability, speech_probability, 0.1, 0.1, 0.1] + assert _drive(config, scores) == [("start", 0), ("stop", 224)] + + +def test_short_pauses_do_not_split_continuous_speech_into_repeated_turns(): + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=500, min_speech_duration_ms=32) + scores = [0.9] + [0.1, 0.9, 0.9, 0.9] * 20 + assert _drive(config, scores) == [("start", 0)] diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py index 5b8130c36..703af6a35 100644 --- a/tests/engine/test_duplex_orchestrator.py +++ b/tests/engine/test_duplex_orchestrator.py @@ -515,3 +515,48 @@ async def test_a_dead_replica_closes_the_sessions_it_was_serving() -> None: messages = [output_q.get_nowait() for _ in range(output_q.qsize())] types = [getattr(getattr(m, "event", None), "type", type(m).__name__) for m in messages] assert types[-1] in {"session.expired", "session.closed"}, f"the client needs a terminal event, got {types}" + + +@pytest.mark.asyncio +async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(monkeypatch): + from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens + from vllm_omni.model_executor.models.qwen3_omni.duplex.plugin import Qwen3OmniDuplexPlugin + + orchestrator, clients, rpc_q, _ = _build() + plugin = Qwen3OmniDuplexPlugin(_encode_audio) + plugin.processor = SimpleNamespace(apply_chat_template=lambda messages, **kwargs: "audio prompt") + orchestrator.plugin = orchestrator.session_manager.plugin = plugin + seen = [] + + def process_inputs(**kwargs): + seen.append(kwargs) + return build_engine_core_request_from_tokens( + request_id=kwargs["request_id"], + prompt={"prompt_token_ids": [1, 2]}, + params=kwargs["params"], + model_config=orchestrator.stage_pools[0].stage_vllm_config.model_config, + resumable=kwargs["resumable"], + ) + + monkeypatch.setattr( + orchestrator, "_get_stage_input_processor", lambda stage_id: SimpleNamespace(process_inputs=process_inputs) + ) + await orchestrator._dispatch_message( + OpenDuplexSessionMessage( + control_id="open-qwen", + session_id=SESSION_ID, + session_config=DuplexSessionConfig(model="qwen", modalities=["text", "audio"]), + ) + ) + assert (await rpc_q.get()).ok + try: + await _submit(orchestrator, _append_audio()) + assert not clients[0].add_request_calls + await _submit(orchestrator, commands.Commit(final=True, create_response=True)) + assert len(seen) == 1 + assert seen[0]["prompt"]["multi_modal_data"]["audio"][0][0].shape == (16000,) + submitted = clients[0].add_request_calls[0][0] + assert submitted.resumable is False + assert not orchestrator.request_states[submitted.request_id].streaming.enabled + finally: + await _close(orchestrator, rpc_q) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index ebc08662f..43d88d692 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -47,6 +47,7 @@ from argparse import Namespace from contextlib import asynccontextmanager from dataclasses import dataclass from types import SimpleNamespace +from typing import Any import pytest from fastapi import FastAPI, HTTPException @@ -564,7 +565,7 @@ async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes start getting stamped / mutated incorrectly. """ - captured: dict[str, object] = {} + captured: dict[str, Any] = {} def fake_build_openai_app(args, supported_tasks): return FastAPI() @@ -592,8 +593,8 @@ async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args()) middleware = captured["served_app"] - http_scope = {"type": "http", "state": {}} - ws_scope = {"type": "websocket"} + http_scope: dict[str, Any] = {"type": "http", "state": {}} + ws_scope: dict[str, Any] = {"type": "websocket"} async def _receive(): return {"type": "http.disconnect"} @@ -920,11 +921,10 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: engine = _FakeEngineClient(stage_configs=[stage]) def _for_diffusion_factory(label: str): - @classmethod def _factory(cls, *args, **kwargs): return _marker(label) - return _factory + return classmethod(_factory) monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) @@ -965,13 +965,11 @@ async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> speech_kwargs = {} def _for_diffusion_factory(label: str): - @classmethod def _factory(cls, *args, **kwargs): return _marker(label) - return _factory + return classmethod(_factory) - @classmethod def _speech_factory(cls, *args, **kwargs): speech_kwargs.update(kwargs) return _marker("speech") @@ -985,7 +983,7 @@ async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> ) monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video")) monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None) - monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _speech_factory) + monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", classmethod(_speech_factory)) monkeypatch.setattr( api_server.ServingRealtimeRobotOpenPI, "create_policy_server", diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py index 9cf58b114..0d9eebe5b 100644 --- a/tests/entrypoints/openai_api/test_duplex_api_server.py +++ b/tests/entrypoints/openai_api/test_duplex_api_server.py @@ -490,3 +490,31 @@ async def test_a_model_that_does_not_declare_chat_completions_does_not_get_the_r assert state.openai_serving_chat is None assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) + + +def test_qwen_plugin_preserves_custom_turn_deployment_default(monkeypatch, tmp_path): + from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **kwargs: QWEN3_OMNI_PIPELINE) + deploy = tmp_path / "qwen.yaml" + deploy.write_text("stages: []\n") + assert not api_server._should_serve_duplex("qwen", {"deploy_config": str(deploy)}) + deploy.write_text("session_mode: duplex\nstages: []\n") + assert api_server._should_serve_duplex("qwen", {"deploy_config": str(deploy)}) + + +@pytest.mark.parametrize("flag", ["1", "true", "on"]) +def test_turn_deployment_rejects_explicit_duplex_without_falling_back_to_stt(flag: str) -> None: + app = _duplex_app(None) + # Even an available STT service must not silently accept a duplex request. + app.state.openai_serving_realtime = object() + with TestClient(app) as client: + with client.websocket_connect(f"/v1/realtime?duplex={flag}") as websocket: + assert websocket.receive_json() == { + "type": "error", + "code": "unsupported", + "error": "VAD realtime is not enabled", + } + with pytest.raises(WebSocketDisconnect) as exc: + websocket.receive_text() + assert exc.value.code == 1008 diff --git a/tests/examples/test_minicpmo_realtime_web_server.py b/tests/examples/test_minicpmo_realtime_web_server.py index 6c35760a7..44f04eb1b 100644 --- a/tests/examples/test_minicpmo_realtime_web_server.py +++ b/tests/examples/test_minicpmo_realtime_web_server.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import hashlib import importlib.util import sys @@ -70,7 +73,12 @@ def test_build_app_injects_public_realtime_url(): def test_build_app_versions_client_bundle_and_disables_index_cache(): client = TestClient(build_app()) expected_version_hash = hashlib.sha256() - for asset_name in ("app.js", "pcm_worklet.js", "playback_worklet.js"): + for asset_name in ( + "app.js", + "pcm_worklet.js", + "playback_worklet.js", + *[str(p.relative_to(STATIC_DIR)) for p in sorted((STATIC_DIR / "profiles").glob("*.js"))], + ): expected_version_hash.update((STATIC_DIR / asset_name).read_bytes()) expected_version = expected_version_hash.hexdigest()[:12] @@ -97,3 +105,41 @@ def test_build_app_exposes_realtime_websocket_and_static_assets(): assert client.get("/static/app.js").status_code == 200 assert client.get("/static/pcm_worklet.js").status_code == 200 assert client.get("/static/playback_worklet.js").status_code == 200 + + +def test_qwen_config_does_not_inject_reference_voice(): + client = TestClient(build_app(profile="qwen3-turn", adapter="vad", model="Qwen/Qwen3-Omni-30B-A3B-Instruct")) + page = client.get("/").text + assert '"profile": "qwen3-turn"' in page + assert '"adapter": "vad"' in page + assert '"refAudio": null' in page + for name in ("common", "minicpm_native", "qwen3_turn"): + assert client.get(f"/static/profiles/{name}.js").status_code == 200 + with pytest.raises(ValueError, match="ref-audio"): + build_app(profile="qwen3-turn", ref_audio="unused.wav") + + +def test_config_cannot_close_its_script_element(): + page = TestClient(build_app(model="")).get("/").text + assert "